Imitation Learning - gymnasium

앞에서 간단히 봤던 Imitation Learning에 대해서, 조금만 더 보도록 하자. 

이전의 간단한 예제는 cartpole과 같이 좌우로 이동하는 것에 대한 예제였다. 

 

이번에는 continuous 한 값을 갖는 예제로 보자. 

차량을 운전하는 모델을 만들어서, 자동차를 운전하도록 시켜보자. 

 

그리고 gymnasium 라이브러리를 가져다 쓰도록 하자. 

 

먼저 학습을 위한 환경을 구성하자. 

import gymnasium as gym
from gymnasium import spaces
import numpy as np

class SimpleCarEnv(gym.Env):
    def __init__(self):
        super(SimpleCarEnv, self).__init__()
        # 상태: [x_position, y_position, heading_angle]
        self.observation_space = spaces.Box(low=np.array([-100, -100, -np.pi]),
                                            high=np.array([100, 100, np.pi]),
                                            dtype=np.float32)
        # 행동: steering_angle (-0.5 rad ~ 0.5 rad)
        self.action_space = spaces.Box(low=np.array([-0.5]),
                                       high=np.array([0.5]),
                                       dtype=np.float32)
        self.target = np.array([10.0, 10.0])
        self.reset()

    def reset(self, seed=None, options=None):
        self.state = np.array([0.0, 0.0, 0.0])
        return self.state, {}

    def step(self, action):
        x, y, theta = self.state
        steering_angle = action[0]

        # simple kinematic model
        velocity = 1.0  # constant speed
        theta += steering_angle
        x += velocity * np.cos(theta)
        y += velocity * np.sin(theta)

        self.state = np.array([x, y, theta])

        # 목표와 거리 계산
        distance = np.linalg.norm(self.target - np.array([x, y]))
        reward = -distance
        terminated = distance < 1.0
        truncated = False

        return self.state, reward, terminated, truncated, {}

    def render(self):
        print(f"State: {self.state}")

 

gym.Env를 상속해서 SimpleCarEnv 클래스를 만든다. 

클래스 내에는 init 과 reset, step 함수를 만드는데, init 에서는 observation_space, action_space를 정의해준다. 

self.target 은 차량이 도착해야하는 목적지를 나타낸다. 

 

reset은 초기화 내용을 넣어주고, 

step 은 action 을 받아서, 해당 action 에 따른 환경의 변화값, 여기서는 차량의 이동위치 및 방향, reward 등을 계산하고 반환한다. 이 반환값을 가지고 다시 차량은 다시 action을 수행하게 된다. 

 

환경은 구성되었으니, 이제 Expert Policy를 만들자. 

Imitation Learning이기 때문에 Expert가 필요하다. 즉 모델이 보고 배울 만한 Expert가 있어야 한다. 간단한 Expert를 만들자. 

def expert_policy(state, target):
    x, y, theta = state
    dx, dy = target[0] - x, target[1] - y
    target_angle = np.arctan2(dy, dx)
    steering = np.clip(target_angle - theta, -0.5, 0.5)
    return np.array([steering], dtype=np.float32)

 

너무 간단히 만들었나? Expert라고 해도 단순히 목표점을 받아서, 그 방향으로 가기 위한 steering 값을 계산하는 것만 한다.

즉, 목표점을 향해 무조건 직진이다. 

 

우리가 보기엔 무조건 직진만 하는 무식한 놈이지만, 아무것도 못하 모델 입장에서보면 목표점까지 갈 수 있는 대단한 Expert 이다. 어찌되었던 목표점까지 가지 않는가? 

 

이제 이 Expert의 행동을 배워서 따라하는 모델을 만들고, Imitation Learning을 해보자. 

import torch
import torch.nn as nn
import torch.optim as optim

# Policy 네트워크 정의
class PolicyNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Sequential(
            nn.Linear(3, 64),
            nn.ReLU(),
            nn.Linear(64, 32),
            nn.ReLU(),
            nn.Linear(32, 1),
            nn.Tanh()  # -1~1 출력
        )

    def forward(self, x):
        return 0.5 * self.fc(x)  # -0.5 ~ 0.5로 스케일

env = SimpleCarEnv()
policy = PolicyNet()
optimizer = optim.Adam(policy.parameters(), lr=0.001)
criterion = nn.MSELoss()

# Expert 데이터 수집
states, actions = [], []
for _ in range(1000):
    state, _ = env.reset()
    for _ in range(50):
        action = expert_policy(state, env.target)
        states.append(state)
        actions.append(action)
        state, _, terminated, _, _ = env.step(action)
        if terminated:
            break

states = torch.tensor(states, dtype=torch.float32)
actions = torch.tensor(actions, dtype=torch.float32)

# imitation learning 학습
for epoch in range(300):
    optimizer.zero_grad()
    outputs = policy(states)
    loss = criterion(outputs, actions)
    loss.backward()
    optimizer.step()

    if (epoch+1) % 50 == 0:
        print(f"Epoch {epoch+1}, Loss: {loss.item():.4f}")

# 테스트
state, _ = env.reset()
for _ in range(20):
    action = policy(torch.tensor([state], dtype=torch.float32)).detach().numpy()[0]
    state, _, terminated, _, _ = env.step(action)
    env.render()
    if terminated:
        print("목표 도달!")
        break

 

작지만 PolicyNet 을 구성하고 forward 함수를 정의하였다. PolicyNet은 모델 자체에서의 최종 출력은 Tanh() 로 -1~1 사이의 출력을 보이고, 이것을 steering 범위 내로 맞추기 위하여 forward 함수에서 0.5배하여 -0.5~0.5 사이의 값을 갖도록 하였다. 

 

그리고 expert의 행동을 수집한다. 

아래 코드 부분인데, state 를 expert_policy 에 입력으로 주고, 그 결과로 action을 받는다. 

이때의 state와 action을 저장하였다가 나중에 딥러닝 모델을 학습시킬 때 데이터로 사용하게 된다. 

즉, Expert가 어떤 환경(state)일 때, 어떤 행동(action)을 수행하는지 기록해두었다가, 그것을 정답으로 모델을 학습시켜서 해당 모델이 Expert와 동일하게 행동하도록 하게 하는 것이다. 

# Expert 데이터 수집
states, actions = [], []
for _ in range(1000):
    state, _ = env.reset()
    for _ in range(50):
        action = expert_policy(state, env.target)
        states.append(state)
        actions.append(action)
        state, _, terminated, _, _ = env.step(action)
        if terminated:
            break

 

실제적인 모델 학습은 아래 코드에서 수행하게 된다. 

학습은 현재 환경(state)를 입력으로 주고, 그에 따른 목표 행동(outputs)을 출력하도록 학습시키며, 그 목표행동(outputs)과 정답에 해당하는 Expert의 행동(action)을 비교하여, 정답 행동과 유사한 행동을 출력하도록 학습시키게 된다. 

states = torch.tensor(states, dtype=torch.float32)
actions = torch.tensor(actions, dtype=torch.float32)

# imitation learning 학습
for epoch in range(300):
    optimizer.zero_grad()
    outputs = policy(states)
    loss = criterion(outputs, actions)
    loss.backward()
    optimizer.step()

    if (epoch+1) % 50 == 0:
        print(f"Epoch {epoch+1}, Loss: {loss.item():.4f}")

 

그리고 아래 코드에서 학습된 PolicyNet을 이용해서 실제로 목표점까지 이동하는지 테스트 한다. 

# 테스트
state, _ = env.reset()
for _ in range(20):
    action = policy(torch.tensor([state], dtype=torch.float32)).detach().numpy()[0]
    state, _, terminated, _, _ = env.step(action)
    env.render()
    if terminated:
        print("목표 도달!")
        break

 

실제로 잘 되는지 실행해보자. 

Imitation Learning을 학습한 모델이 목표까지 잘 도착했다.

 

실제 돌려보면 Expert의 행동을 학습한 모델(PolicyNet)이 목표지점까지 도착한 것을 확인할 수 있다. 

굿! 굿!

 

정리를 해보면, 

현재 환경(state)와 그에 따른 Expert의 행동(action)을 계속해서 저장했다가, 

해당 환경에서 해당 action을 출력하도록 모델을 학습시킨다. 

그렇게 학습된 모델을 이용하여 환경을 입력으로 주면, 그에 맞도록 Expert의 행동(action)과 유사한 행동(output)을 출력으로 낸다. 이것이 Imitation Learning의 가장 간단한 예제가 될 수 있다. 

 

중요한 점은 환경(state)과 행동(action)을 어떻게 정의하느냐의 문제도 있고, 

Imitation Learning의 어려운 점으로 Expert의 행동을 어떻게 수집하느냐의 문제도 있다. 

 

다만 Imitation Learning은 좀 더 발전하여 Reinforcement Learning 으로 나아가는 기초가 될 수 있다. 

Reinforcement Learning 에는 Expert의 행동(action)을 수집하는 대신에, Reward 함수를 이용하여 Expert를 대신한다. 

 

일단은 Expert의 행동을 유사하게 따라하는 방식으로 모델을 학습시켜 문제를 푸는 방법으로 Imitation Learning이 있으며, 

이 구조는 환경에 따른 Expert의 행동을 저장했다가, 그대로 모델을 학습시킨다는 것을 기억하면 좋을 듯 하다. 

 

그럼 오늘은 이만!!