NiiL

RL_W10

Soft Actor-Critic (SAC)

[Soft Actor-Critic (SAC) 알고리즘 개요]

[연속 행동 공간 문제]

[Soft Actor-Critic (SAC) 특징]

1. Gaussian Policy를 활용한 연속적인 행동 표현 및 샘플링

2. 효율적인 Off-policy 학습

3. Q값 과대평가 완화 - 두 개의 Q-네트워크

4. V-네트워크 - Soft State Value 예측

[SAC 네트워크 구조]

1. Q-네트워크 (Q-Network)

(State,Action)↓Input Layer↓Hidden Layer↓Output Layer↓Q\begin{array}{c} \boxed{(State, Action)} \\ \downarrow \\ \boxed{\text{Input Layer}} \\ \downarrow \\ \boxed{\text{Hidden Layer}} \\ \downarrow \\ \boxed{\text{Output Layer}} \\ \downarrow \\ \boxed{Q} \end{array}

2. 타겟 Q-네트워크 (Target Q-Network)

3. 정책 네트워크 (Policy Network)

State↓Input Layer↓Hidden Layer↓Output Layer↓(Meanμ,Stdσ)\begin{array}{c} \boxed{State} \\ \downarrow \\ \boxed{\text{Input Layer}} \\ \downarrow \\ \boxed{\text{Hidden Layer}} \\ \downarrow \\ \boxed{\text{Output Layer}} \\ \downarrow \\ \boxed{(\text{Mean} \;\mu , \text{Std}\;\sigma)} \end{array}

4. V-네트워크 (V-Network)

(State)↓Input Layer↓Hidden Layer↓Output Layer↓V\begin{array}{c} \boxed{(State)} \\ \downarrow \\ \boxed{\text{Input Layer}} \\ \downarrow \\ \boxed{\text{Hidden Layer}} \\ \downarrow \\ \boxed{\text{Output Layer}} \\ \downarrow \\ \boxed{V} \end{array}

[loss 계산 및 update 과정]

1. 샘플링 (Sampling)

2. Target V-value 계산

3. Target Q-value 계산

4. Q-networks Loss (Critic Loss)

5. V-network Loss (Critic Loss)

6. Policy Loss (Actor Loss)

7. Alpha Loss (Temperature Parameter Loss)

8. Network Update

  1. Q-networks 업데이트:
  1. V-network 업데이트:
  1. (Delayed) Policy network 업데이트:
  1. Alpha 업데이트
  1. Target V-network 업데이트:

[실습: Soft Actor-Critic (SAC)를 이용한 Pendulum 실습]

Reference: https://colab.research.google.com/github/MrSyee/pg-is-all-you-need/blob/master/05.SAC.ipynb

1. 라이브러리 가져오기 및 하이퍼파라미터 설정

from typing import Dict, List, Tuple

import gymnasium as gym
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.distributions import Normal
from IPython.display import clear_output

import imageio

# parameters
num_frames = 50000
plotting_interval = 500 # plot refresh 간격 (frames 단위)
memory_size = 100000
batch_size = 256
initial_random_steps = 10000
gamma = 0.99  # 할인율
lr_actor = 0.0003 # learning rate
lr_critic = 0.0003
lr_alpha = 0.0003
tau = 0.005 # target soft update rate 0: no-update 1: hard-update
policy_update_interval = 2
early_stop_threshold = -150 # pendulum-v1 대체로 -200 이상이면 성공으로 분류
early_stop_patience = 10
render_gif_path = "pendulum_simulation_sac.gif"

2. Neural Network 모델 정의

def init_layer_uniform(layer: nn.Linear, init_w: float = 3e-3) -> nn.Linear:
    """Init uniform parameters on the single layer."""
    layer.weight.data.uniform_(-init_w, init_w)
    layer.bias.data.uniform_(-init_w, init_w)

    return layer


class PolicyNetwork(nn.Module):
    def __init__(self, state_dim, action_dim, hidden_dim=256, log_std_min=-20, log_std_max=2):
        super(PolicyNetwork, self).__init__()
        self.log_std_min = log_std_min
        self.log_std_max = log_std_max

        self.linear1 = nn.Linear(state_dim, hidden_dim)
        self.linear2 = nn.Linear(hidden_dim, hidden_dim)

        self.log_std_linear = nn.Linear(hidden_dim, action_dim)
        self.log_std_linear = init_layer_uniform(self.log_std_linear)
        self.mean_linear = nn.Linear(hidden_dim, action_dim)
        self.mean_linear = init_layer_uniform(self.mean_linear)

    def forward(self, state):
        x = F.relu(self.linear1(state))
        x = F.relu(self.linear2(x))

        mean = self.mean_linear(x).tanh()
        log_std = self.log_std_linear(x).tanh()
        log_std = self.log_std_min + 0.5 * (self.log_std_max - self.log_std_min) * (log_std + 1)
        std = torch.exp(log_std)

        # sampling
        dist = Normal(mean, std)
        z = dist.rsample()
        action = z.tanh()

        # Enforce action bounds (for continuous action space) 1e-7: log(0)을 계산하는 상황을 막기 위함
        log_prob = (dist.log_prob(z) - torch.log(1 - action.pow(2) + 1e-7)).sum(axis=-1, keepdim=True)

        return action, log_prob


class QNetwork(nn.Module):
    def __init__(self, state_dim, action_dim, hidden_dim=256):
        super(QNetwork, self).__init__()
        self.linear1 = nn.Linear(state_dim + action_dim, hidden_dim)
        self.linear2 = nn.Linear(hidden_dim, hidden_dim)
        self.linear3 = nn.Linear(hidden_dim, 1)
        self.linear3 = init_layer_uniform(self.linear3)

    def forward(self, state, action):
        x = torch.cat((state, action), dim=-1)
        x = F.relu(self.linear1(x))
        x = F.relu(self.linear2(x))
        x = self.linear3(x)

        return x


class VNetwork(nn.Module):  # Optional: Target V-network
    def __init__(self, state_dim, hidden_dim=256):
        super(VNetwork, self).__init__()
        self.linear1 = nn.Linear(state_dim, hidden_dim)
        self.linear2 = nn.Linear(hidden_dim, hidden_dim)
        self.linear3 = nn.Linear(hidden_dim, 1)
        self.linear3 = init_layer_uniform(self.linear3)

    def forward(self, state):
        x = F.relu(self.linear1(state))
        x = F.relu(self.linear2(x))
        x = self.linear3(x)

        return x

3. Replay Buffer 클래스

class ReplayBuffer:
    """A simple numpy replay buffer."""

    def __init__(self, obs_dim: int, act_dim: int, size: int, batch_size: int = 32):
        """Initializate."""
        self.obs_buf = np.zeros([size, obs_dim], dtype=np.float32)
        self.next_obs_buf = np.zeros([size, obs_dim], dtype=np.float32)
        self.acts_buf = np.zeros([size, act_dim], dtype=np.float32)
        self.rews_buf = np.zeros([size], dtype=np.float32)
        self.done_buf = np.zeros([size], dtype=np.float32)
        self.max_size, self.batch_size = size, batch_size
        self.ptr, self.size = 0, 0

    def store(
        self,
        obs: np.ndarray,
        act: np.ndarray,
        rew: float,
        next_obs: np.ndarray,
        done: bool,
    ):
        """Store the transition in buffer."""
        self.obs_buf[self.ptr] = obs
        self.next_obs_buf[self.ptr] = next_obs
        self.acts_buf[self.ptr] = act
        self.rews_buf[self.ptr] = rew
        self.done_buf[self.ptr] = done
        self.ptr = (self.ptr + 1) % self.max_size
        self.size = min(self.size + 1, self.max_size)

    def sample_batch(self) -> Dict[str, np.ndarray]:
        """Randomly sample a batch of experiences from memory."""
        idxs = np.random.choice(self.size, size=self.batch_size, replace=False)
        return dict(
            obs=self.obs_buf[idxs],
            next_obs=self.next_obs_buf[idxs],
            acts=self.acts_buf[idxs],
            rews=self.rews_buf[idxs],
            done=self.done_buf[idxs],
        )

    def __len__(self) -> int:
        return self.size

3. SAC Agent 클래스

class SACAgent:
    """Simplified SAC agent interacting with environment."""
    def __init__(self, env: gym.Env, seed: int =777):
        """Initialize."""
        obs_dim = env.observation_space.shape[0]
        action_dim = env.action_space.shape[0]

        self.env = env
        self.memory = ReplayBuffer(obs_dim, action_dim, memory_size, batch_size)
        self.seed = seed

        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        print(self.device)

        # Automatic entropy tuning
        self.target_entropy = -np.prod((action_dim,)).item()  # 정해진 규칙은 없으나, 널리 사용되는 방식
        self.log_alpha = torch.zeros(1, requires_grad=True, device=self.device)
        self.alpha_optim = optim.Adam([self.log_alpha], lr=lr_alpha)

        # Networks
        self.actor = PolicyNetwork(obs_dim, action_dim).to(self.device)
        self.v_net = VNetwork(obs_dim).to(self.device)
        self.target_v_net = VNetwork(obs_dim).to(self.device)
        self.target_v_net.load_state_dict(self.v_net.state_dict())
        self.q_net1 = QNetwork(obs_dim, action_dim).to(self.device)
        self.q_net2 = QNetwork(obs_dim, action_dim).to(self.device)

        # Optimizers
        self.actor_optim = optim.Adam(self.actor.parameters(), lr=lr_actor)
        self.v_optim = optim.Adam(self.v_net.parameters(), lr=lr_critic)
        self.q_optim1 = optim.Adam(self.q_net1.parameters(), lr=lr_critic)
        self.q_optim2 = optim.Adam(self.q_net2.parameters(), lr=lr_critic)

        # transition to store in memory
        self.transition = list()

        # total steps count
        self.total_step = 0

        # mode: train / test
        self.is_test = False


    def select_action(self, state: np.ndarray) -> np.ndarray:
        """Select an action from the input state."""
        if self.total_step < initial_random_steps and not self.is_test:
            selected_action = self.env.action_space.sample()
        else:
            act, _= self.actor(torch.FloatTensor(state).to(self.device))
            selected_action = act.detach().cpu().numpy()
        self.transition = [state, selected_action]

        return selected_action

    def step(self, action: np.ndarray) -> Tuple[np.ndarray, float, bool]:
        """Take an action and return the response of the env."""
        # network의 action 출력은 -1~1, pendulum의 action space는 -2~2 이므로 이를 scaling
        action_input = action * self.env.action_space.high
        next_state, reward, terminated, truncated, _ = self.env.step(action_input)
        done = terminated or truncated

        if not self.is_test:
            self.transition += [reward, next_state, done]
            self.memory.store(*self.transition)

        return next_state, reward, done


    def update_model(self) -> None:
        """Update the model by gradient descent."""
        samples = self.memory.sample_batch()
        state = torch.FloatTensor(samples["obs"]).to(self.device)
        next_state = torch.FloatTensor(samples["next_obs"]).to(self.device)
        action = torch.FloatTensor(samples["acts"]).to(self.device)
        reward = torch.FloatTensor(samples["rews"].reshape(-1, 1)).to(self.device)
        done = torch.FloatTensor(samples["done"].reshape(-1, 1)).to(self.device)

        new_action, log_prob = self.actor(state)

        # train alpha (dual problem)
        alpha_loss = (-self.log_alpha.exp() * (log_prob + self.target_entropy).detach()).mean()
        self.alpha_optim.zero_grad()
        alpha_loss.backward()
        self.alpha_optim.step()
        alpha = self.log_alpha.exp()  # used for the actor loss calculation

        # q_net loss
        mask = 1 - done
        q1 = self.q_net1(state, action)
        q2 = self.q_net2(state, action)
        target_v = self.target_v_net(next_state)
        target_q = reward + gamma * target_v * mask
        q_net1_loss = F.mse_loss(q1, target_q.detach())
        q_net2_loss = F.mse_loss(q2, target_q.detach())

        # v_net loss
        v = self.v_net(state)
        q = torch.min(self.q_net1(state, new_action), self.q_net2(state, new_action))
        target_v = q - alpha * log_prob
        v_loss = F.mse_loss(v, target_v.detach())

        if self.total_step % policy_update_interval == 0:
            # actor loss
            advantage = q - v.detach()
            actor_loss = (alpha * log_prob - advantage).mean()

            # actor update
            self.actor_optim.zero_grad()
            actor_loss.backward()
            self.actor_optim.step()

            # target_v_net update
            self._target_soft_update()
        else:
            actor_loss = torch.zeros(())

        # q_net update
        self.q_optim1.zero_grad()
        q_net1_loss.backward()
        self.q_optim1.step()

        self.q_optim2.zero_grad()
        q_net2_loss.backward()
        self.q_optim2.step()

        q_loss = q_net1_loss + q_net2_loss

        # v_net update
        self.v_optim.zero_grad()
        v_loss.backward()
        self.v_optim.step()


    def train(self, num_frames: int = num_frames, plotting_interval: int = 500):
        """Train the agent."""
        self.is_test = False

        state, _ = self.env.reset(seed=self.seed)
        scores = []
        score = 0

        for self.total_step in range(1, num_frames + 1):
            action = self.select_action(state)
            next_state, reward, done = self.step(action)

            state = next_state  # Update current state
            score += reward

            if done:
                state, _ = self.env.reset(seed=self.seed)
                scores.append(score)
                score = 0

            if len(self.memory) >= batch_size and self.total_step > initial_random_steps:
                self.update_model()

            if self.total_step % plotting_interval == 0:
                self._plot(self.total_step, scores)

            # Early stop condition
            if self.total_step > initial_random_steps:
              if np.min(scores[-early_stop_patience:]) >= early_stop_threshold:
                  print(f"Early stopping triggered! at Step: {self.total_step+1} Score: {score}")
                  break

        self.env.close()


    def _target_soft_update(self):
        """Soft-update target network."""
        for target_param, param in zip(self.target_v_net.parameters(), self.v_net.parameters()):
            target_param.data.copy_(tau * param.data + (1.0 - tau) * target_param.data)

    def _plot(self, frame_idx: int, scores: List[float],):
        """Plot the training progresses."""
        clear_output(wait=True)
        plt.figure(figsize=(4, 3))
        plt.title(f"Frame: {frame_idx}, Avg. Score: {np.mean(scores[-10:]):.2f}")
        plt.plot(scores)
        plt.show()

5. 학습 실행 및 결과 시각화

# environment
env = gym.make("Pendulum-v1", render_mode="rgb_array")

# 디버깅을 위해 seed 고정 (같은 seed일 때 같은 연산 유도)
seed = 777
np.random.seed(seed)
if torch.backends.cudnn.enabled:
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True

agent = SACAgent(env, seed)
agent.train(num_frames, plotting_interval)
<Figure size 400x300 with 1 Axes>
실행 결과
Early stopping triggered! at Step: 17401 Score: 0

6. 학습된 에이전트 검증 및 GIF 저장

state, _ = env.reset(seed = agent.seed)
done = False
score = 0
i = 0
frames = []

while not done:
    i = i + 1
    action = agent.select_action(state)
    next_state, reward, done = agent.step(action)
    state = next_state
    score += reward

    # 렌더링 이미지를 프레임 리스트에 추가
    frames.append(env.render())

# GIF 파일로 저장
imageio.mimsave(render_gif_path, frames, duration=33)  # fps 조절 가능 duration = 1000/fps
print(f"Test Episode Score: {score}")

env.close()
Test Episode Score: -127.97587843816552