sac_cartpole.py

import random
import time

import gym
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F

import numpy as np
from copy import deepcopy
from collections import deque


class GumbelSoftmax(torch.distributions.RelaxedOneHotCategorical):
    def sample(self, sample_shape=torch.Size()):
        '''Gumbel-softmax sampling. Note rsample is inherited from RelaxedOneHotCategorical'''
        u = torch.empty(self.logits.size(), device=self.logits.device, dtype=self.logits.dtype).uniform_(0, 1)
        noisy_logits = self.logits - torch.log(-torch.log(u))
        return torch.argmax(noisy_logits, dim=-1)

    def log_prob(self, value):
        '''value is one-hot or relaxed'''
        if value.shape != self.logits.shape:
            value = F.one_hot(value.long(), self.logits.shape[-1]).float()
            assert value.shape == self.logits.shape
        return - torch.sum(- value * F.log_softmax(self.logits, -1), -1)


class SAC(nn.Module):
    def __init__(self,
                 state_dim: int,
                 action_dim: int,
                 pi_lr: float = 1e-3,
                 q_lr: float = 1e-3,
                 gamma: float = 0.99,
                 alpha: float = 1e-3,
                 tau: float = 1e-2,
                 capacity: int = 10_000
                 ):
        super(SAC, self).__init__()
        self.pi_model = nn.Sequential(
            nn.Linear(state_dim, 128),
            nn.ReLU(),
            nn.Linear(128, 128),
            nn.ReLU(),
            nn.Linear(128, action_dim),
        )

        self.q1_model = nn.Sequential(
            nn.Linear(state_dim + action_dim, 128),
            nn.ReLU(),
            nn.Linear(128, 128),
            nn.ReLU(),
            nn.Linear(128, 1)
        )

        self.q2_model = nn.Sequential(
            nn.Linear(state_dim + action_dim, 128),
            nn.ReLU(),
            nn.Linear(128, 128),
            nn.ReLU(),
            nn.Linear(128, 1)
        )

        self.q1_target_model = deepcopy(self.q1_model)
        self.q2_target_model = deepcopy(self.q2_model)

        self.state_dim = state_dim
        self.action_dim = action_dim
        self.gamma = gamma
        self.alpha = alpha
        self.tau = tau
        self.memory = deque([], maxlen=capacity)

        self.pi_optimizer = optim.Adam(self.pi_model.parameters(), lr=pi_lr)
        self.q1_optimizer = optim.Adam(self.q1_model.parameters(), lr=q_lr)
        self.q2_optimizer = optim.Adam(self.q2_model.parameters(), lr=q_lr)

    def predict_action(self, states: torch.FloatTensor):
        logits = self.pi_model(states)
        dist = GumbelSoftmax(temperature=0.01, logits=logits)
        actions = dist.rsample()
        log_probs = dist.log_prob(actions)
        return actions, log_probs

    def get_action(self, state):
        states = torch.FloatTensor(state).unsqueeze(0)
        action, _ = self.predict_action(states)
        return torch.argmax(action).squeeze(0).detach().numpy()

    def update_model(self, loss, optimizer, model=None, target_model=None):
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if (model is not None) and (target_model is not None):
            for param, target_param in zip(model.parameters(), target_model.parameters()):
                new_target_param = (1 - self.tau) * target_param + self.tau * param
                target_param.data.copy_(new_target_param)

    def fit(self, state, action, reward, done, next_state, batch_size: int = 64):
        self.memory.append([state, action, reward, done, next_state])
        if len(self.memory) > batch_size:
            batch = random.sample(self.memory, batch_size)
            states, actions, rewards, dones, next_states = map(np.array, zip(*batch))
            states, actions, rewards, dones, next_states = map(torch.FloatTensor,
                                                               [states, actions, rewards, dones, next_states])
            rewards, dones = rewards.unsqueeze(1), dones.unsqueeze(1)

            next_actions, next_log_probs = self.predict_action(next_states)
            next_states_and_actions = torch.concat([next_states, next_actions], dim=1)
            next_q1_values = self.q1_target_model(next_states_and_actions)
            next_q2_values = self.q1_target_model(next_states_and_actions)
            next_q_values = torch.min(next_q1_values, next_q2_values)
            targets = rewards + self.gamma * (1 - dones) * (next_q_values - self.alpha * next_log_probs)

            actions = F.one_hot(actions.type(torch.LongTensor), self.action_dim)

            states_and_actions = torch.concat([states, actions], dim=1)
            q1_loss = torch.mean((targets.detach() - self.q1_model(states_and_actions)) ** 2)
            q2_loss = torch.mean((targets.detach() - self.q2_model(states_and_actions)) ** 2)
            self.update_model(q1_loss, self.q1_optimizer, self.q1_model, self.q1_target_model)
            self.update_model(q2_loss, self.q2_optimizer, self.q2_model, self.q2_target_model)

            pred_actions, pred_log_probs = self.predict_action(states)
            states_and_pred_actions = torch.concat([states, pred_actions], dim=1)
            pred_q1_values = self.q1_target_model(states_and_pred_actions)
            pred_q2_values = self.q1_target_model(states_and_pred_actions)
            pred_q_values = torch.min(pred_q1_values, pred_q2_values)
            pi_loss = - torch.mean(pred_q_values - self.alpha * pred_log_probs)
            self.update_model(pi_loss, self.pi_optimizer)


def visualize(env, agent, max_len=1000):
    trajectory = {'states': [], 'actions': [], 'rewards': []}

    obs = env.reset()
    state = obs

    for _ in range(max_len):
        trajectory['states'].append(state)

        action = agent.get_action(state)
        trajectory['actions'].append(action)

        obs, reward, done, _ = env.step(action)
        trajectory['rewards'].append(reward)

        state = obs

        time.sleep(0.03)
        env.render()

        if done:
            break

    return trajectory


if __name__ == '__main__':
    env = gym.make('CartPole-v1')
    env.reset()
    state_dim = env.observation_space.shape[0]
    action_dim = env.action_space.n

    print(f"n_states: {state_dim}, n_actions: {action_dim}")

    episode_n = 100
    session_len = 500
    total_rewards = []

    agent = SAC(state_dim, action_dim)

    for episode in range(episode_n):
        state = env.reset()
        total_reward = 0
        for _ in range(session_len):
            action = agent.get_action(state)
            next_state, reward, done, _ = env.step(action)
            agent.fit(state, action, reward, done, next_state)

            total_reward += reward
            state = next_state

        total_rewards.append(total_reward)

        print(f'episode {episode}, reward: {total_reward}')

    visualize(env, agent)