05_AC.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Actor-Critic Methods\n",
    "\n",
    "In _value-based_ methods, we use a value function to determine the optimal policy. The exploration\n",
    "and exploitation tradeoff is left to manual tuning. In _policy-based_ methods, we learn the policy\n",
    "directly, but the Monte Carlo methods have high-variance and tend to be slower (mitigating that by\n",
    "collecting more samples causes the algorithm to be less sample-efficient).\n",
    "\n",
    "Actor-critic methods combine the two approaches:\n",
    "\n",
    " * We learn the _actor_ policy, to control how the agent behave.\n",
    " * We measure how good the chosen actions are via a _critic_, which learns the value function.\n",
    "\n",
    "Remember the _baseline_ in policy gradient? The _critic_ helps computing it!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import torch.optim as optim\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "import gymnasium as gym\n",
    "\n",
    "from util.gymnastics import DEVICE, ReplayBuffer, gym_simulation\n",
    "from util.gymnastics import init_random, plot_scores, soft_update_model_params"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gym_simulation(\"Pendulum-v1\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# For convenience, harcoding actions' interval [-2.0, 2.0]\n",
    "ACTION_SCALE = 2.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## General (1-Step) Algorithm\n",
    "\n",
    "The simplified version of an [actor-critic](https://arxiv.org/abs/1602.01783) algorithm goes as\n",
    "follows:\n",
    "\n",
    " 1. In state $S_t$, the _actor_ outputs the action $A_t$ using the policy $\\pi_{\\theta}(S_t)$ and\n",
    "    obtaining $R_{t+1}$ and $S_{t+1}$\n",
    " 2. The _critic_ outputs the value of state and next state: $V_t = \\^{v_w}(S_t),\n",
    "    V_{t+1} = \\^{v_w}(S_{t+1})$\n",
    " 3. We compute the _advantage_: $A_t = Q(s_t, a_t) - V(s_t) = [R_{t+1} + \\gamma V_{t+1}] - V_t$\n",
    " 3. The _actor_ updates the policy parameters $\\theta$ using the _advantage_:\n",
    "    $\\space \\space \\Delta \\theta = \\alpha \\nabla_{\\theta} [ log \\pi_{\\theta}(S_t) ] A_t$\n",
    " 4. The _critic_ updates the value function parameters $w$ minimizing $A_t$\n",
    " 5. Repeat\n",
    "\n",
    "The _advantage actor critic (A2C)_ algorithm works basically this way. For a robust implementation,\n",
    "check [Stable Baselines3 A2C](https://stable-baselines3.readthedocs.io/en/master/modules/a2c.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Neural Network Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ActorNetwork(nn.Module):\n",
    "    \"\"\"Actor (Policy) Model.\"\"\"\n",
    "\n",
    "    def __init__(self, state_size, action_size, fc1_units=400, fc2_units=300):\n",
    "        super(ActorNetwork, self).__init__()\n",
    "        # TODO: Create three linear layers, first two initialized via Kaiming normal, the last one\n",
    "        #       with uniform distribution in [-3e-3, 3e-3]. the input of the policy network is a\n",
    "        #       state, while the output a (continuous) action.\n",
    "\n",
    "    def forward(self, state):\n",
    "        \"\"\"Build an actor (policy) network that maps states -> actions.\"\"\"\n",
    "        # TODO: Use ReLU, ReLU, TanH :)\n",
    "        pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CriticNetwork(nn.Module):\n",
    "    \"\"\"Critic (Value) Model.\"\"\"\n",
    "\n",
    "    def __init__(self, state_size, action_size, fcs1_units=400, fc2_units=300):\n",
    "        super(CriticNetwork, self).__init__()\n",
    "        # TODO: Same architecture as the policy network, but now the input is the flattened state\n",
    "        #       and action, while the output is a single value: the state-action value.\n",
    "\n",
    "    def forward(self, state, action):\n",
    "        \"\"\"Build a critic (value) network that maps (state, action) pairs -> Q-values.\"\"\"\n",
    "        # TODO: concat state and action (pay attention at the dimension!), ReLU for non-linearity.\n",
    "        #       The output is directly the output of the last linear layer (without non-linearity).\n",
    "        pass\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Q-Learning Based Actor-Critic Algorithms\n",
    "\n",
    "We are going to implement modern actor-critic algorithms such as DDPG, TD3, SAC. They are closely\n",
    "related to Q-Learning, and in-between DQN and policy-gradient methods. All these algorithms are\n",
    "analyzed in detail in [OpenAI SpinningUp](https://spinningup.openai.com/en/latest) and I strongly\n",
    "recommend to read that amazing resource!\n",
    "\n",
    "That is because these algorithms learn approximators for the optimal action-value function\n",
    "$Q^*(s, a)$ and optimal (deterministic, excluding SAC) policy $a^*(s) = \\argmax_{a} Q^*(s, a)$\n",
    "exploting the fact that the action-value function is differentiable.\n",
    "\n",
    "Finding the max of the action-value function would be an optimization problem in and of itself, but\n",
    "instead we learn a deterministic policy $\\mu(s)$ such that: $Q(s, \\mu(s)) \\approx \\max_a Q^(s,a)$.\n",
    "All of these algorithms are off-policy and use a replay buffer (like DQN).\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training Loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_actor_critic(env, agent, n_episodes=1_000, max_t=300):\n",
    "    \"\"\"General training loop for actor-critic algorithms in this lecture.\"\"\"\n",
    "    # Records all episode scores.\n",
    "    scores = []\n",
    "    for i_episode in range(1, n_episodes+1):\n",
    "        # TODO: Reset the environment.\n",
    "        state, _ = None\n",
    "        # TODO: Reset the score to zero.\n",
    "        score = None\n",
    "        # Run training for max_timesteps.\n",
    "        for _ in range(max_t):\n",
    "            # TODO: Select an action via agent.act(state, add_noise=True) -> notice the noise!\n",
    "            action = None\n",
    "            # TODO: Perform a step in the environment.\n",
    "            next_state, reward, terminated, truncated, _ = None\n",
    "            # TODO: Compute `done`.\n",
    "            done = None\n",
    "            # TODO: Perform a step for the agent calling agent.step(...)\n",
    "            # ...\n",
    "            # Update state, score, and check termination.\n",
    "            state = next_state\n",
    "            score += reward\n",
    "            if done:\n",
    "                break\n",
    "        # Record statistics and print debugging information.\n",
    "        scores.append(score)\n",
    "        avg_score = np.mean(scores[-100:])\n",
    "        print(f'\\rEpisode {i_episode}\\tAverage Score: {avg_score:.2f}',\n",
    "              end=\"\\n\" if i_episode % 50 == 0 else \"\")\n",
    "        if avg_score >= -370:\n",
    "            print(f'\\rEpisode {i_episode} solved environment!\\tAverage Score: {avg_score:.2f}')\n",
    "            break\n",
    "            \n",
    "    return scores"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## DDPG\n",
    "\n",
    "The _Deep Deterministic Policy Gradient_ algorithm can be thought the DQN algorithm for continuous\n",
    "action spaces, and it uses the same techniques: replay-buffer + target-networks (for both actor and\n",
    "critic).\n",
    "\n",
    "We already leant how the _actor_ learns a deterministic policy $\\mu(s)$ maximizing the critic value.\n",
    "The _critic_ learns the action-value function $Q(s, a)$ minimizing a _mean-squared Bellman error_\n",
    "(MSBE - the squared _advantage_ if you will) like DQN.\n",
    "\n",
    "Because the policy is deterministic, the _exploration / exploitation_ tradeoff can be tuned adding\n",
    "noise to the action. In the [original paper](https://proceedings.mlr.press/v32/silver14.pdf), the\n",
    "authors recommended to use\n",
    "[Ornstein-Uhlenbeck noise](https://en.wikipedia.org/wiki/Ornstein%E2%80%93Uhlenbeck_process), but\n",
    "it turns out that a zero-mean Gaussian noise works just as well (hence, we'll use that here)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AgentDDPG:\n",
    "    def __init__(self, state_size, action_size, start_mem_size=128,\n",
    "                 gamma=0.99, lr_actor=1e-4, lr_critic=1e-3, exploration_noise_scale=0.1):\n",
    "        self.state_size = state_size\n",
    "        self.action_size = action_size\n",
    "        self.start_mem_size = start_mem_size\n",
    "        self.gamma = gamma\n",
    "        self.exploration_noise_scale = exploration_noise_scale\n",
    "\n",
    "        # Actor network (w/ target network)\n",
    "        # TODO: Create actor local and target networks, plus the Adam optimizer. Make sure to keep\n",
    "        #       the target network in eval mode and copy the local network parameters.\n",
    "        self.actor        = None\n",
    "        self.actor_target = None\n",
    "        # ...\n",
    "        self.actor_optimizer = None\n",
    "\n",
    "        # Critic network (w/ target network)\n",
    "        # TODO: Create critic local and target networks, plus the Adam optimizer. Make sure to keep\n",
    "        #       the target network in eval mode and copy the local network parameters.\n",
    "        self.critic        = None\n",
    "        self.critic_target = None\n",
    "        # ...\n",
    "        self.critic_optimizer = None\n",
    "\n",
    "        # TODO: Create the replay buffer (using the provided one in the `util` module).\n",
    "        self.memory = None\n",
    "\n",
    "    @torch.no_grad\n",
    "    def act(self, state: np.array, add_noise=False):\n",
    "        \"\"\"Returns actions for given state as per current policy.\"\"\"\n",
    "        # TODO: Convert the state into a tensor.\n",
    "        state = None\n",
    "        # TODO: Get the action via the actor local network\n",
    "        action = None\n",
    "        # In the original paper, the noise is generated via an Ornstein-Uhlenbeck process. It turns\n",
    "        # out, a normal gaussian noise works just as well. Hence, that's what we use.\n",
    "        if add_noise:\n",
    "            # TODO: Add noise via normal distribution (np.random.normal) times the parameter\n",
    "            #       `self.exploration_noise_scale`.\n",
    "            action += None\n",
    "        # TODO: Return the action, making sure to clip it based on ACTION_SCALE.\n",
    "        return None\n",
    "\n",
    "    def step(self, state, action, reward, next_state, done):\n",
    "        \"\"\"Save experience in replay memory, and use random sample from buffer to learn.\"\"\"\n",
    "        # TODO: Save experience / reward in the memory.\n",
    "        # ...\n",
    "\n",
    "        # TODO: Learn, if enough samples are available in memory.\n",
    "        if len(self.memory) > self.start_mem_size:\n",
    "            pass\n",
    "\n",
    "    def learn(self, experiences):\n",
    "        states, actions, rewards, next_states, dones = experiences\n",
    "\n",
    "        # CRITIC UPDATE.\n",
    "        with torch.no_grad():\n",
    "            # TODO: Get the next actions using the actor target on next_states.\n",
    "            actions_next = None\n",
    "            # TODO: Get the Q_targets_next using the critic target.\n",
    "            Q_targets_next = None\n",
    "            # TODO: Compute Q targets for current states: rewards + (gamma * Q_targets * (1-dones))\n",
    "            Q_targets = None\n",
    "\n",
    "        # TODO: Compute the Q_expected using the critic local network.\n",
    "        Q_values = None\n",
    "        # TODO: Compute critic loss: MSE between Q_expected and Q_targets.\n",
    "        critic_loss = None\n",
    "        # TODO: Perform a minimization step of the critic loss with its optimizer.\n",
    "        # ...\n",
    "\n",
    "        # ACTOR UPDATE.\n",
    "        # TODO: Compute the action predictions via the actor_local network.\n",
    "        actions_pred = None\n",
    "        # TODO: Compute actor loss, which is the negative critic_local(states, actions_pred) mean.\n",
    "        actor_loss = None\n",
    "        # TODO: Perform a minimization step of the actor loss with its optimizer.\n",
    "        # ...\n",
    "\n",
    "        # TODO: update target networks, calling the `soft_update_model_params` utility function.\n",
    "        # ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with init_random(gym.make('Pendulum-v1')) as env:\n",
    "    agent_ddpg = AgentDDPG(env.observation_space.shape[0], env.action_space.shape[0])\n",
    "    scores_ddpg = train_actor_critic(env, agent_ddpg)\n",
    "plot_scores(scores_ddpg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gym_simulation(\"Pendulum-v1\", agent_ddpg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## TD3\n",
    "\n",
    "The _Twin Delayed DDPG_ (TD3) algorithm expands on DDPG with a couple of additional tricks:\n",
    "\n",
    " * It learns _two_ Q functions, and uses the smaller Q value to for the target. That is to address\n",
    "   overestimation of Q values in DDPG. The \"_twin_\" part of the name comes from this.\n",
    " * Updates the policy (and target) networks less frequently than the Q function (hence, \"_delayed_\")\n",
    "   and that is to keep the target and learning more stable.\n",
    " * Finally, it adds noise to the target action to \"smooth out\" the action value and make it harder\n",
    "   for the policy to exploit errors in the Q function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AgentTD3:\n",
    "    def __init__(self, state_size, action_size, start_mem_size=128,\n",
    "                 gamma=0.99, lr_actor=1e-4, lr_critic=1e-3, exploration_noise_scale=0.1,\n",
    "                 policy_noise = 0.2, noise_clamp=0.5,\n",
    "                 policy_freq=2):\n",
    "        self.state_size = state_size\n",
    "        self.action_size = action_size\n",
    "        self.start_mem_size = start_mem_size\n",
    "        self.gamma = gamma\n",
    "        self.exploration_noise_scale = exploration_noise_scale\n",
    "        self.policy_noise = policy_noise\n",
    "        self.noise_clamp = noise_clamp\n",
    "        self.policy_freq = policy_freq\n",
    "        self.t_step = 0\n",
    "\n",
    "        # Actor network (w/ target network)\n",
    "        # TODO: Build the same actor local and target network + optimizer as DDPG.\n",
    "        # ...\n",
    "\n",
    "        # TD3 trick n.1: Twin critic networks (w/ target network)\n",
    "        # TODO: Build **two** twin critic networks!\n",
    "        self.twin_critic_1        = None\n",
    "        # ...\n",
    "        self.twin_critic_2        = None\n",
    "        # ...\n",
    "\n",
    "        # TODO: Build a single Adam optimizer (hint: concatenate all parameters as list).\n",
    "        self.critic_optimizer     = None\n",
    "\n",
    "        # TODO: Instantiate the replay buffer.\n",
    "        self.memory = ReplayBuffer()\n",
    "\n",
    "    @torch.no_grad\n",
    "    def act(self, state, add_noise=False):\n",
    "        \"\"\"Returns actions for given state as per current policy.\"\"\"\n",
    "        # TODO: Convert the state into a tensor.\n",
    "        state = None\n",
    "        # TODO: Get the action via the actor local network\n",
    "        action = None\n",
    "        # In the original paper, the noise is generated via an Ornstein-Uhlenbeck process. It turns\n",
    "        # out, a normal gaussian noise works just as well. Hence, that's what we use.\n",
    "        if add_noise:\n",
    "            # TODO: Add noise via normal distribution (np.random.normal) times the parameter\n",
    "            #       `self.exploration_noise_scale`.\n",
    "            action += None\n",
    "        # TODO: Return the action, making sure to clip it based on ACTION_SCALE.\n",
    "        return None\n",
    "\n",
    "    def step(self, state, action, reward, next_state, done):\n",
    "        \"\"\"Save experience in replay memory, and use random sample from buffer to learn.\"\"\"\n",
    "        self.t_step += 1\n",
    "\n",
    "        # TODO: Save experience / reward in the memory.\n",
    "        # ...\n",
    "\n",
    "        # TODO: Learn, if enough samples are available in memory.\n",
    "        if len(self.memory) > self.start_mem_size:\n",
    "            pass\n",
    "\n",
    "    def learn(self, experiences):\n",
    "        states, actions, rewards, next_states, dones = experiences\n",
    "\n",
    "        # UPDATE TWIN CRITICS.\n",
    "        with torch.no_grad():\n",
    "            # TODO: Get the next actions using the actor target.\n",
    "            actions_next = None\n",
    "\n",
    "            # TD3 trick n.3: target policy smoothing.\n",
    "            # TODO: Get the noise torch.randn_like(actions_next) times the policy_noise.\n",
    "            noise = None\n",
    "            # TODO: Clamp the noise based on noise_clamp and scale it by ACTION_SCALE\n",
    "            noise = None\n",
    "            # TODO: Add the noise to the actions_next.\n",
    "            actions_next += None\n",
    "            # TODO: Clamp the actions to be in the correct interval [-ACTION_SCALE, ACTION_SCALE].\n",
    "            actions_next = None\n",
    "\n",
    "            # TODO: Compute Q_targets_1 and Q_targets_2 via the twin critic networks using the next\n",
    "            #       states and actions.\n",
    "            Q_targets_next_1 = None\n",
    "            Q_targets_next_2 = None\n",
    "            # TODO: Pick Q_targets_next as the minimum between those two.\n",
    "            Q_targets_next = None\n",
    "            # TODO: Compute Q_targets as DDPG :)\n",
    "            Q_targets = None\n",
    "\n",
    "        # TODO: Compute the Q_values of both twin critics.\n",
    "        Q_values_1 = None\n",
    "        Q_values_2 = None\n",
    "        # TODO: Compute the critic loss as sum of mse_loss of the two critics.\n",
    "        critic_loss = None\n",
    "\n",
    "        # TODO: Minimize the critic loss.\n",
    "        # ...\n",
    "\n",
    "        # UPDATE ACTOR.\n",
    "        # TD3 trick n.2: delayed policy updates.\n",
    "        if self.t_step % self.policy_freq == 0:\n",
    "            # TODO: Get actions from the actor.\n",
    "            actions_pred = None\n",
    "            # TODO: Compute the Q_target via the first twin critic.\n",
    "            Q_values = None\n",
    "            # TODO: Compute the action loss as negative mean of the Q_target.\n",
    "            actor_loss = None\n",
    "\n",
    "            # TODO: Minimize the actor loss.\n",
    "            # ...\n",
    "\n",
    "            # TODO: Update all target networks.\n",
    "            # ...\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with init_random(gym.make('Pendulum-v1')) as env:\n",
    "    agent_td3 = AgentTD3(env.observation_space.shape[0], env.action_space.shape[0])\n",
    "    scores_td3 = train_actor_critic(env, agent_td3)\n",
    "plot_scores(scores_td3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gym_simulation(\"Pendulum-v1\", agent_td3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SAC (Optional)\n",
    "\n",
    "The _Soft Actor Critic_ (SAC) algorithm is an off-policy algorithm similar to DDPG and TD3, that\n",
    "learns a stochastic policy instead. It adopts many of the techniques used in TD3, but it stems from\n",
    "the _Maximum Entropy Formulation_ of reinforcement learning.\n",
    "\n",
    "For an in-depth understanding of both max-ent and SAC, I suggest watching\n",
    "[Lecture 1](https://www.youtube.com/watch?v=2GwBez0D20A&list=PLwRJQ4m4UJjNymuBM9RdmB3Z9N5-0IlY0&index=2)\n",
    "of Peter Abbeel Deep RL course, and reading the OpenAI SpinningUp SAC summary.\n",
    "\n",
    "In short and as a highlight, in this formulation of RL the optimization objective is to maximize the\n",
    "expected return plus the _entropy_ of the policy $H[ \\pi (a | s_t)]$, term that intuitively balance\n",
    "exploration and exploitation:\n",
    "\n",
    "$$\n",
    "\\max_{\\pi} \\mathbb{E}\\Bigl[ \\sum_{t=0}^{T} r_t + \\beta H[ \\pi (a | s_t)] \\Bigr]\n",
    "$$\n",
    "\n",
    "The _entropy_ in fact \"measures\" how uncertain is a policy (i.e., a deterministic policy has a very\n",
    "low entropy, while a random one has high entropy).\n",
    "\n",
    "While the _critic_ network learns the action-value as usual, the _actor_ network learns the mean\n",
    "and standard deviation of a Gaussian representing the stochastic policy, using `tanh` to \"squash\"\n",
    "the values in an acceptable range, and having the `logstd` depending on the network parameters as\n",
    "well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ActorNetworkSAC(nn.Module):\n",
    "    \"\"\"The actor network for SAC. It is provided given its technicalities.\"\"\"\n",
    "\n",
    "    def __init__(self, state_size, action_size, action_scale=ACTION_SCALE, action_bias=0.0):\n",
    "        super().__init__()\n",
    "        self.fc1 = nn.Linear(state_size, 256)\n",
    "        self.fc2 = nn.Linear(256, 256)\n",
    "        self.fc_mean = nn.Linear(256, action_size)\n",
    "        self.fc_logstd = nn.Linear(256, action_size)\n",
    "        self.action_scale = action_scale\n",
    "        self.action_bias = action_bias\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.fc1(x)\n",
    "        x = F.relu(x)\n",
    "        x = self.fc2(x)\n",
    "        x = F.relu(x)\n",
    "        mean    = self.fc_mean(x)\n",
    "        log_std = self.fc_logstd(x)\n",
    "        log_std = torch.tanh(log_std)\n",
    "        return mean, self.adjust_log_std(log_std)\n",
    "\n",
    "    def get_action(self, x):\n",
    "        mean, log_std = self(x)\n",
    "        std = log_std.exp()\n",
    "        normal = torch.distributions.Normal(mean, std)\n",
    "        sample = normal.rsample() # Reparameterization trick: (mean + std * N(0,1))\n",
    "        output = torch.tanh(sample)\n",
    "        log_prob = normal.log_prob(sample)\n",
    "        # Enforcing action bounds (and non-zero log)\n",
    "        action = output * self.action_scale + self.action_bias\n",
    "        log_prob -= torch.log(self.action_scale * (1 - output.pow(2)) + 1e-6)\n",
    "        return action, log_prob\n",
    "\n",
    "    def adjust_log_std(self, log_std):\n",
    "        log_std_min, log_std_max = (-5, 2) # From SpinUp / Denis Yarats\n",
    "        return log_std_min + 0.5 * (log_std_max - log_std_min) * (log_std + 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AgentSAC:\n",
    "    def __init__(self, state_size, action_size, start_mem_size=128,\n",
    "                 gamma=0.99, lr_actor=1e-4, lr_critic=1e-3, policy_freq=2):\n",
    "        self.state_size = state_size\n",
    "        self.action_size = action_size\n",
    "        self.start_mem_size = start_mem_size\n",
    "        self.gamma = gamma\n",
    "        self.policy_freq = policy_freq\n",
    "        self.t_step = 0\n",
    "\n",
    "        # TODO: Build the actor network with the ActorNetworkSAC (no target!)\n",
    "        self.actor = None\n",
    "        self.actor_optimizer = optim.Adam(self.actor.parameters(), lr=lr_actor)\n",
    "\n",
    "        # TODO: Build twin critics like in TD3.\n",
    "        # ...\n",
    "\n",
    "        # TODO: Make the replay buffer.\n",
    "        self.memory = None\n",
    "\n",
    "    @torch.no_grad\n",
    "    def act(self, state, add_noise=False): # SAC doesn't really have noise, but for consistency...\n",
    "        \"\"\"Returns actions for given state as per current policy.\"\"\"\n",
    "        # TODO: Convert the numpy state to a tensor.\n",
    "        state = None\n",
    "        # TODO: Get the action calling `get_action` from the actor.\n",
    "        action, _ = None\n",
    "        # TODO: Return the numpy action.\n",
    "        return None\n",
    "\n",
    "    def step(self, state, action, reward, next_state, done):\n",
    "        \"\"\"Save experience in replay memory, and use random sample from buffer to learn.\"\"\"\n",
    "        self.t_step += 1\n",
    "\n",
    "        # TODO: Save experience / reward in the memory.\n",
    "        # ...\n",
    "\n",
    "        # TODO: Learn, if enough samples are available in memory.\n",
    "        if len(self.memory) > self.start_mem_size:\n",
    "            pass\n",
    "\n",
    "    def learn(self, experiences, alpha=0.2):\n",
    "        states, actions, rewards, next_states, dones = experiences\n",
    "\n",
    "        # UPDATE TWIN CRITICS\n",
    "        with torch.no_grad():\n",
    "            # TODO: Get next action and logprob via actor.get_action(states).\n",
    "            actions_next, log_pi_next_st = None\n",
    "            # TODO: Compute the entropy term as: alpha * log_pi_next_st\n",
    "            entropy_term = None\n",
    "\n",
    "            # TODO: Compute Q_targets like in TD3, BUT subtract entropy_term from Q_targets_next.\n",
    "            # ...\n",
    "\n",
    "        # TODO: Compute and minimize critic loss like in TD3.\n",
    "        # ...\n",
    "\n",
    "        # UPDATE ACTOR.\n",
    "        if self.t_step % self.policy_freq == 0:\n",
    "            # TODO: Get action and log via actor.get_action(...)\n",
    "            action, log_pi = None\n",
    "            # TODO: Compute the entropy term as above\n",
    "            entropy_term = None\n",
    "\n",
    "            # TODO: Use the min of the twin critic computed Q values.\n",
    "            Q_values_1 = None\n",
    "            Q_values_2 = None\n",
    "            Q_values = None\n",
    "            # TODO: actor loss is the (entropy_term - Q_values).mean()\n",
    "            actor_loss = None\n",
    "\n",
    "            # TODO: Minimize the actor loss.\n",
    "            # ...\n",
    "\n",
    "            # TODO: Update the twin critic target networks.\n",
    "            # ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with init_random(gym.make('Pendulum-v1')) as env:\n",
    "    agent_sac = AgentSAC(env.observation_space.shape[0], env.action_space.shape[0])\n",
    "    scores_sac = train_actor_critic(env, agent_sac)\n",
    "plot_scores(scores_sac)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gym_simulation(\"Pendulum-v1\", agent_sac)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "drlzh",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}