📌 JAX로 강화 학습 모델 구축 - Q-learning을 이용한 CartPole 문제 해결

📌 JAX로 강화 학습 모델 구축 - Q-learning을 이용한 CartPole 문제 해결

---

🚀 강화 학습이란?

강화 학습(RL)은 **에이전트(Agent)**가 환경(Environment)과 상호작용하며, 보상(Reward)을 최대화하는 정책(Policy)을 학습하는 기법입니다.
이번 글에서는 JAX를 활용하여 Q-learning 기반의 강화 학습 모델을 구축하여 CartPole 문제를 해결해보겠습니다.

---

💡 1. Q-learning의 핵심 개념

📐 Q-learning 알고리즘 구조

• Q 함수: 상태-행동 쌍 (s, a)의 가치를 추정하는 함수
• 목표: Q값을 업데이트하여 최적 정책을 학습
• Bellman Equation:

Q(s,a)=Q(s,a)+α×[r+γ×max⁡a′Q(s′,a′)−Q(s,a)]Q(s, a) = Q(s, a) + \alpha \times [r + \gamma \times \max_a' Q(s', a') - Q(s, a)]

• α\alpha: 학습률 (Learning Rate)
• γ\gamma: 할인율 (Discount Factor)
• rr: 보상

---

🔧 2. 환경 설정

📥 필수 라이브러리 설치

pip install gym
pip install jax jaxlib
pip install optax  # JAX 최적화 라이브러리

🗺️ 환경 불러오기 (Gym)

반응형

import gym

# CartPole 환경 설정
env = gym.make("CartPole-v1")
state = env.reset()
print(f"초기 상태: {state}")

---

📝 3. Q 네트워크 모델 정의

🧠 신경망 구조

• 입력: 상태 벡터 (4차원)
• 출력: 행동 가치 (2차원: 좌/우 이동)
• 은닉층: 128개의 노드, ReLU 활성화
• 출력층: 상태-행동 가치

import jax
import jax.numpy as jnp
from jax import random
import optax  # JAX용 최적화 라이브러리

# 네트워크 초기화 함수
def init_params(key, input_dim, hidden_dim, output_dim):
    key1, key2 = random.split(key)
    w1 = random.normal(key1, (input_dim, hidden_dim)) * 0.1
    b1 = jnp.zeros(hidden_dim)
    w2 = random.normal(key2, (hidden_dim, output_dim)) * 0.1
    b2 = jnp.zeros(output_dim)
    return (w1, b1), (w2, b2)

# Q 네트워크
def q_network(params, x):
    (w1, b1), (w2, b2) = params
    hidden = jnp.tanh(jnp.dot(x, w1) + b1)  # ReLU 대신 Tanh 사용
    q_values = jnp.dot(hidden, w2) + b2
    return q_values

---

📉 4. 손실 함수와 최적화

🧮 손실 함수 (Mean Squared Error)

def loss_fn(params, state, action, target):
    q_values = q_network(params, state)
    q_value = q_values[action]
    return jnp.mean((target - q_value) ** 2)

🔧 옵티마이저 설정

optimizer = optax.adam(learning_rate=0.001)

# 최적화 상태 초기화
@jax.jit
def init_optim_state(params):
    return optimizer.init(params)

# 파라미터 업데이트 함수
@jax.jit
def update(params, opt_state, state, action, target):
    loss, grads = jax.value_and_grad(loss_fn)(params, state, action, target)
    updates, opt_state = optimizer.update(grads, opt_state)
    new_params = optax.apply_updates(params, updates)
    return new_params, opt_state, loss

---

🔁 5. 학습 루프 정의

🌟 하이퍼파라미터 설정

epsilon = 0.1
gamma = 0.99
episodes = 300
batch_size = 32
buffer_size = 10000
learning_starts = 1000
target_update_freq = 50

🗃️ 경험 리플레이 버퍼

import collections
import random

# 경험 버퍼 정의
buffer = collections.deque(maxlen=buffer_size)

def store_transition(state, action, reward, next_state, done):
    buffer.append((state, action, reward, next_state, done))

def sample_batch(batch_size):
    batch = random.sample(buffer, batch_size)
    states, actions, rewards, next_states, dones = zip(*batch)
    return jnp.array(states), jnp.array(actions), jnp.array(rewards), jnp.array(next_states), jnp.array(dones)

---

🏃 6. 모델 학습

🚀 학습 루프

key = random.PRNGKey(42)
params = init_params(key, 4, 128, 2)
opt_state = init_optim_state(params)

for episode in range(episodes):
    state = env.reset()
    episode_reward = 0

    for t in range(200):
        # Epsilon-greedy 정책
        if random.random() < epsilon:
            action = env.action_space.sample()
        else:
            q_values = q_network(params, jnp.array(state))
            action = int(jnp.argmax(q_values))

        # 행동 수행
        next_state, reward, done, _ = env.step(action)
        store_transition(state, action, reward, next_state, done)

        # 학습 시작 조건
        if len(buffer) > learning_starts:
            states, actions, rewards, next_states, dones = sample_batch(batch_size)

            # Q 타겟 계산
            next_q_values = q_network(params, next_states)
            targets = rewards + gamma * jnp.max(next_q_values, axis=1) * (1 - dones)

            # 파라미터 업데이트
            params, opt_state, loss = update(params, opt_state, states, actions, targets)

        state = next_state
        episode_reward += reward

        if done:
            break

    print(f"Episode {episode + 1}, Reward: {episode_reward}")

    # 타겟 네트워크 업데이트
    if episode % target_update_freq == 0:
        target_params = params

---

✅ 7. 테스트와 성능 평가

def evaluate(env, params, episodes=10):
    total_reward = 0
    for _ in range(episodes):
        state = env.reset()
        episode_reward = 0
        done = False
        while not done:
            q_values = q_network(params, jnp.array(state))
            action = int(jnp.argmax(q_values))
            state, reward, done, _ = env.step(action)
            episode_reward += reward
        total_reward += episode_reward
    return total_reward / episodes

avg_reward = evaluate(env, params)
print(f"평균 테스트 보상: {avg_reward}")

---

📌 다음 글 예고: JAX로 Transformer 모델 구현하기

다음 글에서는 JAX를 활용하여 Transformer 모델을 구현하고,
자연어 처리(NLP) 문제를 해결하는 데 적용해보겠습니다.

---

 

JAX, 강화 학습, Q-learning, CartPole, 딥러닝, Python, GPU 학습, 에이전트 학습, 고속 연산, 모델 학습, 데이터 전처리, 강화 학습 모델, 고성능 컴퓨팅, Gym 환경, 모델 평가