NiiL

RL_W10

English translation

Soft Actor-Critic (SAC)

[Soft Actor-Critic (SAC) algorithm overview]

[The continuous action space problem]

[Soft Actor-Critic (SAC) features]

1. Represent and sample continuous actions using a Gaussian policy

2. Efficient off-policy learning

3. Reduce Q-value overestimation — two Q-networks

4. V-network — Predict the soft state value

[SAC network architecture]

1. 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. 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-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 calculation and update process]

1. Sampling

2. Compute the target V-value

3. Compute target Q-values

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-network updates:
  1. Update the V-network:
  1. (Delayed) policy network update:
  1. Alpha update
  1. Update the target V-network:

[Exercise: Pendulum with Soft Actor-Critic (SAC)]

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

1. Import libraries and set hyperparameters

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 interval (in frames)
memory_size = 100000
batch_size = 256
initial_random_steps = 10000
gamma = 0.99  # Discount factor
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: Generally classified as successful at -200 or higher
early_stop_patience = 10
render_gif_path = "pendulum_simulation_sac.gif"

2. Define the neural network model

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: Prevent evaluating 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

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

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()  # A widely used approach, though no fixed rule exists
        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."""
        # Scale the network action output from -1~1 to the pendulum action space of -2~2
        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. Run training and visualize results

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

# Fix seed for debugging (use the same seed to reproduce the same computations)
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. Validate the trained agent and save a 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

    # Append the rendered image to the frame list
    frames.append(env.render())

# Save as a GIF file
imageio.mimsave(render_gif_path, frames, duration=33)  # Adjust fps as needed: duration = 1000/fps
print(f"Test Episode Score: {score}")

env.close()
Test Episode Score: -127.97587843816552