2023. 12. 13. 17:33ㆍ카테고리 없음
안녕하세요. 오늘은 강화학습의 Policy Gradient 알고리즘 중 가장 기초가 되는 알고리즘인 REINFORCE에 대해서 공부해보겠습니다.
이 글을 통해 REINFORCE 알고리즘의 원리와 어떻게 구현하는지를 아실 수 있습니다. 구현 코드는 Gymansium 공식문서에 있는 예제코드를 초심자들이 쉽게 접근할 수 있도록 수정하였습니다.
재미있는 사실은 REINFROCE는 " ‘RE’ward ‘I’ncrement ‘N’on-negative ‘F’actor times ‘O’ffset ‘R’einforcement times ‘C’haracteristic ‘E’ligibility "의 줄임말이라고 합니다.
심층강화학습에서는 크게 두 가지 방식의 알고리즘이 있습니다.
1. Value based (ex. DQN)
2. Policy based (ex. REINFORCE, DDPG)
DQN은 continuous state space에는 적용이 가능했지만, continuous action space에서는 적용이 불가능했습니다. 그 이유는 무수히 많은 수의 Q값을 구하는 것이 불가능하기 때문입니다. 하지만 Policy를 바로 학습하는 REINFORCE알고리즘의 network는 action의 확률분포를 출력하기 때문에 continuous action space에도 적용이 가능한 방법입니다.
예를 들어, 구현할 때 사용할 InvertedPendulum-v4에서는 action space가 [-3N ,3N]의 힘이 됩니다. Policy network에서 action의 평균과, 표준편차값을 얻어내고, 이 값을 이용해 만든 확률분포에서 action을 고르게 됩니다.
REINFORCE를 사용하면 Experience replay와 target network를 사용하지 않아도됩니다.
Experience replay를 사용했던 이유는 temporal correlation때문이었는데, REINFORCE에서는 하나혹은 여러개의 에피소드로 학습을 하기 때문에 사용할 필요가 없습니다. 그리고 REINFORCE는 MC기반이기 때문에 target network도 필요없습니다. 그저 total reward를 최대로 만드는 방향으로 업데이트하면 됩니다.
Policy Gradient에서는 objective function $J(\theta) = \mathbb{E}_{\pi_\theta} [r(\tau)]$의 gradient를 구해주어야합니다. 증명과정은 다음과 같습니다.
더 궁금하신 분은 오승상 교수님의 강화학습 강의를 참고하시길 바랍니다.
증명과정에서 눈여겨볼것은 Markov property가 적용되었다는 것입니다.
우리가 사용할 공식은 다음과 같습니다.
$$ \textbf{[Policy Gradient Theorem]} \\
\nabla_{\theta} J(\theta) = \nabla_{\theta} \mathbb{E}_{\pi_\theta} [r(\tau)] = \mathbb{E}_{\pi_\theta} \left[ r(\tau) \sum_{t=0}^{T-1} \nabla_{\theta} \log \pi_\theta (a_t | s_t) \right] $$
이제 구현으로 넘어가 코드를 설명해드리겠습니다.
환경분석
https://gymnasium.farama.org/environments/mujoco/inverted_pendulum/
구현에 앞서 학습시킬 환경에 대해서 알아야합니다.
Action Space는 [-3.0N, 3.0N]의 범위를 가지는 연속공간이며, InvertedPendulum에 가해지는 토크값이 됩니다. size는 (1,)입니다.
Observation space는 InvertedPendulum의 상태를 나타내며 연속공간에 해당합니다. numpy배열로 반환되며, 순서대로 카트의 위치, 막대의 수직각도, 카트의 선속도, 막대의 각속도를 나타냅니다.
Reward는 환경안에서 설정한 특정각도안에 머물면 매 timestep마다 +1의 보상이 주어집니다.
Episode는 1000번의 timestep이 넘어가면 truncation이 True가 되면서 종료되고, 막대의 각도가 0.2rad이상이되거나, 상태값중 어느하나라도 유한한 값이 아니라면 terminated가 True가 되면서 종료됩니다.
Box는 gymnasium에서 제공하는 여러 공간중 한개입니다. 환경이 이산적인 값을 가지는 경우에는 Discrete를 사용합니다. (ex. Cartpole-v1)
main 함수에 대해서 알아봅시다.
env = gym.make("InvertedPendulum-v4")
total_num_episodes = int(5e3)
obs_space_dims = env.observation_space.shape[0] # observation space의 차원을 알아내는 방법
action_space_dims = env.action_space.shape[0] # action space의 차원을 알아내는 방법
for seed in [1,2,3,5]:
torch.manual_seed(seed)
random.seed(seed)
np.random.seed(seed)
agent = REINFORCE(obs_space_dims, action_space_dims)
# 출력 설정
print_interval = 100
score = 0
for episode in range(total_num_episodes):
# gymnasium v26 requires users to set seed while resetting the environment
obs, info = env.reset(seed=seed)
done = False
while not done:
action = agent.sample_action(obs)
obs, reward, terminated, truncated, info = env.step(action)
agent.rewards.append(reward)
score += reward
done = terminated or truncated
agent.update()
if episode % print_interval == 0 and episode != 0:
print("seed : %d, episode = %d, avg_score : %.2f" %(seed, episode, score/print_interval))
score = 0
위에서 seed를 여러개로 하는 이유는 seed별로 성능이 왔다갔다하는데 이를 평균내서 일관된 성능을 알아내기 위함입니다. episode에서 action을 뽑고 다음 state로 갈 때 랜덤성이 추가되므로 각각의 랜덤함수에 seed를 통해 일관된 랜덤값이 나올 수 있도록 설정해줍니다.
done이 True가 되기까지 계속해서 환경을 탐험합니다. reward값은 update시에 필요하므로 agent의 rewards변수에 저장해둡니다. 그리고 에피소드가 끝나면 update를 진행합니다.
Policy_Network의 구현에 대해서 알아봅시다.
Fig-1을 보면 shared_net이 있고, 평균, 표준편차를 출력하는 각각의 네트워크가 있습니다. 이를 pytorch를 통해 구현해줍니다. forward함수는 네트워크의 순전파를 담당하는 함수입니다.
표준편차를 구할 때 torch.log를 사용하는 이유는 네트워크의 출력이 혹시나 음수라도 음수 안되도록 해주는 연산입니다.
이렇게 하면 학습 안정성도 올라간다고 하네요.
class Policy_Network(nn.Module):
"""Parametrized Policy Network."""
def __init__(self, obs_space_dims: int, action_space_dims: int): # invpend obs_dim = 4, action_dim = 1
super().__init__()
hidden_space1 = 16 # Nothing special with 16, feel free to change
hidden_space2 = 32 # Nothing special with 32, feel free to change
# Shared Network
self.shared_net = nn.Sequential(
nn.Linear(obs_space_dims, hidden_space1),
nn.Tanh(),
nn.Linear(hidden_space1, hidden_space2),
nn.Tanh(),
)
# Policy Mean specific Linear Layer
self.policy_mean_net = nn.Sequential(
nn.Linear(hidden_space2, action_space_dims)
)
# Policy Std Dev specific Linear Layer
self.policy_stddev_net = nn.Sequential(
nn.Linear(hidden_space2, action_space_dims)
)
def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
shared_features = self.shared_net(x.float())
action_means = self.policy_mean_net(shared_features)
action_stddevs = torch.log(
1 + torch.exp(self.policy_stddev_net(shared_features))
)
return action_means, action_stddevs
REINFORCE 객체에는 에이전트에 대한 정보가 들어있습니다. sample_action함수에서는 action을 고른 후에 업데이트에 사용될 확률의 로그값을 self.probs에 저장해둡니다. 그리고 step함수에서는 tensor가 아니라 numpy를 전달해야하므로 numpy로 변환 후에 action을 리턴합니다. update는 따로 밑에서 설명하겠습니다.
class REINFORCE:
def __init__(self, obs_space_dims: int, action_space_dims: int):
# Hyperparameters
self.learning_rate = 1e-4 # Learning rate for policy optimization
self.gamma = 0.99 # Discount factor
self.eps = 1e-6 # small number for mathematical stability
self.probs = [] # Stores probability values of the sampled action
self.rewards = [] # Stores the corresponding rewards
self.net = Policy_Network(obs_space_dims, action_space_dims)
self.optimizer = torch.optim.AdamW(self.net.parameters(), lr=self.learning_rate)
def sample_action(self, state: np.ndarray) -> float:
state = torch.tensor(np.array([state]))
action_means, action_stddevs = self.net(state)
distrib = Normal(action_means[0] + self.eps, action_stddevs[0] + self.eps)
action = distrib.sample()
prob = distrib.log_prob(action)
action = action.numpy()
self.probs.append(prob)
return action
def update(self):
running_g = 0
gs = []
for R in self.rewards[::-1]:
running_g = R + self.gamma * running_g
gs.insert(0, running_g)
deltas = torch.tensor(gs)
loss = 0
for log_prob, delta in zip(self.probs, deltas):
loss += log_prob.mean() * delta * (-1)
# Update the policy network
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
# Empty / zero out all episode-centric/related variables
self.probs = []
self.rewards = []
Policy Gradient theroem을 유도할때는 total reward를 사용했지만, 구현할때에는 리턴을 사용합니다.
gs는 각 스텝에서의 리턴 값들을 저장해놓은 리스트입니다. 리턴을 뒤에서부터 계산하면 쉽게 계산할 수 있습니다.
$G_{t-2}$의 값은 $t-2$ 이후의 모든 보상의 합인데 이는 $G_{t-1}$에다가 $t-1$의 보상을 합하기만 하면 되기때문입니다.
이 코드에서는 이를 [::-1]을 통해 간단하게 구현하였습니다.
list 슬라이싱의 문법을 사용한 것으로 기본구조는 [start:stop:step]인데 start, step은 생략해서 디폴트로 처음과 끝을 가리키고, step은 -1로 역순으로 출력하는 방법입니다.
my_list = [1, 2, 3, 4, 5]
reversed_list = my_list[::-1] # [5, 4, 3, 2, 1]
그리고 가장 중요한 backward를 위한 loss함수를 정의하는 것입니다.
Pseudo코드에서는 하나하나씩 업데이트를 하지만 실제 코드에서는 한꺼번에 처리합니다. 하나하나씩도 처리하는 코드를 구현해봤는데 "RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation" 다음과 같은 오류가 뜨며 학습이 되질 않았습니다. 따라서 모든 loss를 더한후에 한번에 업데이트를 해줘야합니다.
def update(self):
running_g = 0
gs = []
for R in self.rewards[::-1]:
running_g = R + self.gamma * running_g
gs.insert(0, running_g)
deltas = torch.tensor(gs)
loss = 0
for log_prob, delta in zip(self.probs, deltas):
loss += log_prob.mean() * delta * (-1)
# Update the policy network
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
# Empty / zero out all episode-centric/related variables
self.probs = []
self.rewards = []
<reference>
[1] : https://gymnasium.farama.org/tutorials/training_agents/reinforce_invpend_gym_v26/