Soft Actor Critic (SAC) is a model-free Reinforcement Learning (RL) method developed by Tuomas Haarnoja in a series of papers culminating in Soft Actor-Critic Algorithms and Applications. It is one of the seminal RL papers for continuous control and has inspired much of my work. In this blog post, I’ll review the theory and derivation behind SAC from zero to a fully-derived algorithm without skipping (almost) any of the mathematical details. This time I haven’t included an implementation of my own, but a good one to go along with this blog post is denisyarats/pytorch_sac.

Note: If you are familiar with RL basics, feel free to skip ahead

Table of contents

• Problem
• Bellman Equations
• TD Learning
• Q-Learning
• Policy Gradients
• Actor critic methods
• Soft Actor Critic
• Soft Actor Critic with automatic temperature
• Refernces
• Further resources
• Acknowledgements

Problem

We want a robot with limited perceptive abilities to learn to achieve a given task. Let’s formulate this mathematically and ground it in the real world via a beer-opening robot:

• The world has a state at a given time $s_t \in S$. This state space can consist of anything meaningful to the robot: the position of the beer bottle? how full is it? how thirsty is the human we need to serve?
• Unfortunately, robots can’t open beer bottles by just thinking about it really hard; instead, they have to act in the world via actions $a_t \in A$. For our robot, actions could be a 4-dimensional torque vector to control all robot arm joints.
• To incentivise the robot to achieve its goal, we give it a reward for each timestep $r_t \in \mathbb{R}$. This could be a function of the state and action $r_t := r(s_t, a_t)$ or something more abstract like our human telling us how happy they are. The critical assumption is that the robot only gets a scalar reward $r_t$ at each timestep without knowing how it is calculated.

The goal of our agent now is to maximize the cumulative reward by choosing the optimal sequence of actions:

Technically our reward can depend on several factors from previous timesteps (e.g. did the human think I look cool while opening their beer). However, that quickly becomes an exponentially more complex problem, especially theoretically. To make the problem more tractable and our lives easier, we will make an assumption:

Assumption 1: Given a sequence of state-action pairs $(s_t, a_t) \ \ \ \forall t \in [0, T]$, the transition distribution is fully defined by $p(s_{t+1} | s_t, a_t)$. This is known as the Markov Assumption.

In English: the last state $s_t$ contains all the useful information for our problem. From this we automatically get:

Corollary 1: The reward is defined as a function of our state and action $r: S \times A \rightarrow \mathbb{R}$

Visually this can be look as the following Markov Chain:

With this assumption we can now redefine our objective function as

Still, this is incorrectly defined. Why? Because we defined our transition dynamics as a probabilistic distribution in Assumption 1. To tackle that, we have to define our objective above as an expectation over the rewards:

Since we’re already dealing with probabilistic distributions and ultimately want our results to fit in modern machine learning, we can constrain our actions to be coming from a set of possible policies:

Assumption 2: Actions are sampled from a probabilistic feedback policy $\pi$ parametarised by $\theta \in \mathbb{R}^d$ and conditioned only on the current state $s_t$:

Thus our final (and I mean it this time) objective becomes

Bellman Equations

These problems get really messy really fast. Luckily we have a simple unified way to think about these problems - via the lens of value functions:

Definition 1: The value of a state $s_t$ following policy $\pi$ is given by

English: This is the reward we expect to receive continuing from our current state $s_t$. We will refer to this as the value function.

Let’s expand this and see where we end up

This recursive relationship is known as the Bellman Equation and will be key to our future derivations.

There is also another way to look at our value function from Definition 1 through the lens of an already taken action:

Definition 2: The value of a state-action pair $(s_t, a_t)$ following policy $\pi$ is given by:

We will refer to this as the state-action value function or Q function.

Similarly to value function, here we also have a recursive relationship:

Additionally note the relationship between $V(s_t)$ and $Q(s_t, a_t)$:

TD Learning

Cool, but what can we use these value functions for? Let’s consider the problem of traversing a graph with the highest possible cumulative reward. We don’t select actions; we only transition according to the probabilities in the graph below. Reward is $r(s_6) = 0$ and $r(s) = -1 \ \ \ \forall s \neq s_6$.

One way to solve this is via Dynamic Programming (DP). You can read more about it in my other blog post about Differential Dynamic Programming. What we need to know for now, though, is that we simply need to apply the Bellman Equation (Eq. 3) recursively to all states going from left to right:

Thus, by following the path of highest values, we can conclude that the best path is $s_2 \rightarrow s_4 \rightarrow s_6$. This is the power of value functions and how they allow us to estimate optimality and make decisions (well, in this example, we don’t actually decide).

We say that Dynamic Programming (DP):

• does full backups as it considers all possible transitions in Eq. 3. This works only for small problems.
• bootstraps as it uses other value estimates to update its own new estimates.

If we don’t want to bootstrap, we can instead use Monte Carlo (MC) samples to update our value estimates based on the full return of some sampled trajectories. In other words, we take multiple samples of Eq. 3 and use them to update our $V(s)$ estimates for all visited states $s$.

The full return in our toy example involves reaching the last termination state. We call this an episode.

There are multiple ways to MC estimates. The easiest one is first-visit MC, where we use the return following the first visit.

Using the same toy graph as above, say we are given 4 episodes:

1. $s_1 \rightarrow s_3 \rightarrow s_3 \rightarrow s_6$
   • results in $V(s_1) = -3$, $V(s_3)=-2$, $V(s_6)=0$
2. $s_2 \rightarrow s_4 \rightarrow s_6$
   • results in $V(s_2) = -2$, $V(s_4) = -1$, $V(s_6) = 0$
3. $s_2 \rightarrow s_5 \rightarrow s_6$
   • results in $V(s_2) = -2$, $V(s_5) = -1$, $V(s_6) = 0$
4. $s_2 \rightarrow s_5 \rightarrow s_5 \rightarrow s_6$
   • results in $V(s_2) = -3$, $V(s_5) = -2$, $V(s_6) = 0$

If we average out all the samples we get:

Comparing these values to what we obtained in DP, we can see that they aren’t entirely correct but are close. If we have infinite MC samples, we will get the DP results according to the Weak Law of Large Numbers (WLLN). However, even with these 4 samples, we can still conclude that $s_2 \rightarrow s_4 \rightarrow s_6$ is the optimal path.

What is quite cool about MC is that we don’t need a model (i.e. it is model-free)! The not-so-cool thing is that we need to wait for episodes to finish before we can update our value estimates.

If we don’t want to wait until the end, we can use 1-step MC samples and bootstrap the rest:

where $\alpha$ is an update step and $\hat{V}(s_t)$ is our new estiamte for $V(s_t)$ (that we’re not fully certain of; that’s why we call it estimate). OK, but how can we get an estimate? We can use the definition of the value function from Eq. 2, but we must wait for the episode’s outcome. In Eq. 3, we discovered an interesting recursive relationship that only considers one step in the future. Swapping in that, we get a new update rule:

Thus, now we have an update rule that we just need to wait for 1 timestep before applying! Note that we removed all remnants of $\pi$ and $a_t$ since we aren’t executing actions in this toy example.

This is called Temporal Difference (TD) Learning, and because we use only a single-step sample, this particular update is called a TD(0) update. This is a cross between DP and MC trying to harness the benefits of both. We can frequently update our estimates like DP by bootstrapping but also do so in a model-free fashion like MC!

TD-Learning is considered by many the core of RL! Armed with Eq. 6, we can conceive a simple value estimator algorithm for our toy problem.

Let’s apply this to our toy problem. Initialise $V(s) = 0 \ \ \forall s$, $\alpha = 0.5$.

1. $s_1 \rightarrow s_3 \rightarrow s_3 \rightarrow s_6$
   • $V(s_1) = V(s_1) + 0.5(-1+V(s_3)-V(s_1)) = -0.5$
   • $V(s_3) = V(s_3) + 0.5(-1+V(s_3)-V(s_3)) = -0.5$
   • $V(s_3) = V(s_3) + 0.5(-1+V(s_6)-V(s_3)) = -0.75$
   • $V(s_6) = V(s_6) = 0$
2. $s_2 \rightarrow s_4 \rightarrow s_6$
   • $V(s_2) = V(s_2) + 0.5(-1 + V(s_4) - V(s_2)) = -0.5$
   • $V(s_4) = V(s_4) + 0.5(-1 + V(s_6) - V(s_4)) = -0.5$
   • $V(s_6) = V(s_6) = 0$
3. $s_2 \rightarrow s_5 \rightarrow s_6$
   • $V(s_2) = V(s_2) + 0.5(-1 + V(s_5) - V(s_2)) = -0.75$
   • $V(s_5) = V(s_5) + 0.5(-1 + V(s_6) - V(s_5)) = -0.5$
   • $V(s_6) = V(s_6) = 0$
4. $s_2 \rightarrow s_5 \rightarrow s_5 \rightarrow s_6$
   • $V(s_2) = V(s_2) + 0.5(-1 + V(s_5) - V(s_2)) = -1.125$
   • $V(s_5) = V(s_5) + 0.5(-1 + V(s_5) - V(s_5)) = -1$
   • $V(s_5) = V(s_5) + 0.5(-1 + V(s_6) - V(s_5)) = -1$
   • $V(s_6) = V(s_6) = 0$

Oh wow, we got the almost optimal path again! Let’s compare all of our value estimates from the different methods

TD(0) didn’t do too badly! With more samples and more tuned $\alpha$ it can do even better! (Hint: try $\alpha=0.8$)

Q-Learning

TD-Learning is nice, but in the toy example above, we didn’t have a policy and thus an opportunity to take action. If we had the option, could we use TD-Learning to learn a policy?

One way is to take our TD(0) update from Eq. 6 and reformulate it with state-action values instead:

OK, how do we select actions? The most naive way is by iterating over all actions to find the one with the highest value. Assuming a discrete and finite action space $A = \mathbb{N}^m$, then it is feasible to compute this.

Let’s for a moment say we can feasibly compute this, then why should we compute the expected value of $Q(s_{t+1}, a_{t+1})$ in Eq. 7, why don’t we directly optimize Q there, which is guaranteed to lead to the optimal Q-value? Wait, what is the optimal Q-value?

Definition 3: The optimal value function $V^*(s)$ or $Q^*(s,a)$ is the value of following the optimal policy $\pi^*$ that achieves the highest possible reward on our task: $\forall s \in S, \forall \pi \neq \pi^*$ the following inequality holds $V^*(s) \geq V^\pi(s)$

Thus we can now reformulate Eq. 7 to always follow the (currently) optimal action, not just blindly use the sampled action:

For completeness and references to other literature, using Eq. 6 to update value estimates is known as SARSA, and using Eq. 8 is known as Q-Learning. Both are forms of TD-Learning. These two update rules represent on-policy and off-policy learning, respectively, which are two different schools in the RL community. On-policy updates have nice theoretical guarantees and, in practice, result in faster convergence but require us to learn from the current policy interacting with the environment. In contrast, off-policy approaches allow us to learn from any data and have better exploration but worse theoretical convergence guarantees.

OK, let’s get to something practical - how do we use Q-learning to solve interesting tasks?

When Q is parameterised with a neural net, we arrive at the first successful deep RL algorithm - Deep Q-Network (DQN). In that paper, the authors parameterise the Q-value neural net with some parameters $\psi$ and train it to the loss function

which results in the update rule

Okay, but where do we get the data to train on $(s_t, a_t, r_t, s_{t+1})$? With TD(0), we can learn from 1-step interactions, so we can do step, learn, step, learn and etc… However, NNs aren’t great at learning from a single data point; matter of fact, they only work when trained with Stochastic Gradient Descent (SGD). To make NNs work for Q-Learning, Mnih et al. introduce the concept of an experience replay buffer - a dataset containing the M latest interactions. If M is large enough (e.g. $10^6$), we can sample mini-batches from it and train our Q-network. Thus, we arrive at the famous DQN algorithm!

Policy Gradients

Q-learning is cool and works for deterministic action spaces, but let us now return to continuous action spaces $a_t \in A = R^d$. Finding $\max_a Q(s_t, a)$ becomes infeasibile. Even worse, as we increase $d$, we quickly run into the curse of dimensionality.

Let’s completely disregard Q-learning and return to having a parameterised policy $\pi(a, s; \theta)$ and optimise it to try to solve our base objective from Eq. 1. We can explicitly frame this as

Theorem 1: Policy Gradients (for finite horizon)

If our starting state follows some distribution $p(s_0)$, the gradient of the objective in Eq. 9 is

Pretty cool eh? This tells us that to optimize our objective, we don’t need to differentiate neither the reward, nor the model! This keeps up with our initial assumptions that we have no information about them.

Proof:

The key to the proof is the log-prob trick using the simple calculus identity

If we apply this directly to Eq. 9 we get

Now, this is useful only if finding $\nabla_\theta \log J(\theta)$ is easy. Let’s check that

Where $p(s_0 | s_{-1}, a_{-1}) = 1$

Let’s now take the log of the whole expression above

Now taking the gradient of the whole expression, we see that most terms go away as they aren’t a function of $\theta$

Plugging this directly into Eq. 11

Credits: I pieced together this proof from Sutton’s Chapter 13 and Tedrake’s Chapter 20.

Note: There are many other flavours of proofs for this theorem. One that is interesting to figure out is the infinite time horizon one.

We just developed this theory of optimising our objective; what now? Can we apply it to our problems and solve them? Kind of… The easiest approach we can derive is the Reinforce algorithm (Sutton 1999).

Assuming a finite-horizon problem, we can estimate take episodic Monte-Carlo samples of episode returns (objective Eq. 9) and use our Policy Gradients theorem to optimise some parametarised policy $\pi_\theta(a_t | s_t)$. To ease notation, we define a return sample as: $R_t = \sum_{i=t}^T r(s_i, a_i)$. Then the Reinforce algorithm is as follows

1. Collect episode {$s_0, a_0, r_0, …, r_{T-1}, s$} following policy $\pi_\theta(\cdot | s_t)$
2. For $ t = 0, …, T$
   • $\theta \leftarrow \theta + \alpha R_t \nabla_\theta \log \pi_\theta (a_t | s_t)$

Reinforce is quite powerful but also inconviniently (and inefficiently) requires us to wait for an episode to finish before we can update our policy. This is especially bad when we are dealing with infinite horizon problems where $T = \infty$. Policy gradients (PG) methods are very useful, but they can’t figure out the long-term results of their actions - i.e. we just optimised to state $s_T$ but is that a good state to be in? It may be in the short term, but it may not be in the long term. I love giving the example of the drunken spider who had too much beer at the pub and is now trying to get home.

He is drunk, so his movement is uncertain (i.e. stochastic), and his planning horizon is severely limited. As a matter of fact, he is so drunk that he can only plan as much as points A and B. If we leave the spider to decide based on PG with a simple reward like the closest distance to home, it will decide that point A is better. However, then the spider has to cross the bridge, and if he stumbles, he will fall into the lake, which has a pretty bad reward. Thus, getting to state B is better, but PG needs more foresight to decide this due to their limited time horizon.

Actor critic methods

In a nutshell, actor-critic methods address the limitations of PG (lack of foresight) and Q-learning (inability to scale to continuous action spaces) by combining both ideas into one uniform framework called actor-critic methods. In a nutshell, they have three major components:

1. A $\psi$-parametarised critic that lears a $Q(s,a; \psi)$ state-action value estimator based on Q-learning. We use our standard Eq. 8 but now we get our optimal action from $a_t \sim \pi(a_t | s_t; \theta)$. Thus our loss and update equations are now:

1. A $\theta$-parametarised actor which learns according to our PG theorem Eq. 10. The key here is that instead of optimizing for maximal reward, we can optimize for the surrogate of the reward, the Q-function!

1. Just like DQN, as we are using NNs here, it is a good idea to have an experience replay buffer and train both the actor and the critic back to back with sampled mini-batches and SGD.

Putting all these 3 blocks together we arrive at another famous algorithm called Deep Determinisitc Policy Gradients (DDPG) first presented by Lillicrap et al. in 2015

DDPG is pretty cool as it can now bring all the benefits of Q-learning into continuous action spaces, and it was the first method to successfully do more complex locomotion tasks. For the purposes of this blog post, we will investigate the (half-)cheetah and humanoid tasks of the Deepmind control suite:

DDPG does quite well to the cheetah task but not so well on the humnoid one. Now (finally) we get to the main ordeal…

Soft Actor Critic

DDPG works well with some tricks added in, but it tends to be unstable and lacks exploration. This results in a tendency to get stuck in local minima without solving the task.

To alleviate this, Haarnoja et al. suggest we go back to using stochastic policies (unlike DDPG) and view the problem through the lens of entropy maximisation. For this, we look at entropy.

Definition 4: The entropy of a random variable $X$ designates the information value of $X$:

The intuition here is that less probable samples have a higher value(i.e. entropy) than more probable ones. Let’s say we have a distribution:

Computing the entropy, we get that $X=b$ has higher entropy:

That’s the intuition for a single sample, but the intuition for the entropy of the RV itself is counter intuitive. It measures how random a random variable is! In other words, how well distributed it is. Let’s look at two examples

Cool, what now?

Entropy enables us to open another school of ideas: Maximum Entropy Reinforcement Learning (MaxEnt RL). In this line of work we redefine our initial objective in Eq. 1 as

where $\alpha$ is just a weighing coefficient (called temperature). In English, this reads as maximise the reward while being as random as possible. Somewhat counterintuitive, but this allows us to explore more efficiently and thus find a solution to the task. A new $J(\theta)$ means we also get new value functions:

$V^\pi(s_t)$ is nice but remeber that we need to decide on actions, so a state-action value is more useful:

We get the last line by plugging in Eq. 14. This is An interesting result! The Q function recursive relationship is different from the standard one we see in Q-learning (Eq. 5). The result in Eq. 16 is so important that it has a name: Soft Q-function. Inevitably this would lead us to Soft Actor Critic. First, let’s finish the critic part. With Eq. 16 and a $Q(s_t, a_t; \psi)$ parameterised by a neural net, we can formulate a new regression target and update it with SGD, just like we did in DQN:

which results in the update rule

That handles the critic; now, how do we learn a policy? Since our policy is stochastic, that is non-trivial. I’m not sure how this part was derived in the SAC paper, but to me, it looks reverse engineered (please email me if you have a better clue!).

Our fundamental goal is to maximise the objective in Eq. 12. It’s clear to see that can be translated to something in terms of the Q-function:

Thus we’d expect the gradient to be something proportional to $\log \pi (a | s;\theta) - Q(s, a)$, but unfortunately, we can’t just differentiate when $\pi$ is stochastic. We can be smarter with stochastic calculus and use the KL divergence:

Using the definition of KL divergence we can expand this:

Sampling data from a replay buffer and using our PG Theorem, we can again optimize this with SGD:

This update equation looks nice, but $\nabla_\theta \pi$ isn’t straightforward. To make this tractable, we can restrict the policy to a set of well-known distributions like Gaussians. Then to make this backpropagatable, we can employ the reparametarization trick. If se set our policy to a diagonal Gaussian $\pi(a_t |s_t;\theta) = \mathcal{N}(\mu_\theta, \sigma_\theta)$, we can reparametarise it with a unit diagonal Gaussian as $\pi(a_t | s_t;\theta) = \mu_\theta + \sigma_\theta^T \mathcal{N}(0, 1)$. Then our new parameterised update rule becomes:

TODO expand on reparameterisation?

With this, we can now complete our first basic SAC algorithm:

When applied to the cheetah problem, SAC does well but can’t even match DDPG in terms of both training or evaluation.

SAC does well initially and is more stable than DDPG, but as we approach the optimal policy, performance starts getting worse. The issue is that fixed-temperature SAC has one major limitation, as the name implies - the $\alpha$ parameter is fixed! This means that during training, we care equally about exploration. This is why SAC matches DDPG at the early stages of training but later starts lagging behind.

Naturally, though, we want to explore the early stages of training. Once we’ve done some exploration, we want to prioritise solving the task (i.e. maximising the reward) as efficiently as possible. The straightforward to achieve this would be to decrease $\alpha \rightarrow 0$ over time. That would probably work, but remember that we’re using stochastic policies. They can tell us how confident they are about their actions! Thus we can use the confidence of the policy itself to adjust $\alpha$ automatically!

Soft Actor Critic with automatic temperature

In their 2018 paper, Haarnoja et al. extend SAC to automatically adjust its entropy temperature $\alpha$ by targeting a minimum entropy level.

They do this by reframing the maximum entropy RL objective in Eq. 12 to maintain a minimum entropy level at all times instead of explicitly maximising entropy. Mathematically this translates to the constrained objective

Where $H_0$ is the minimum entropy we are trying to maintain. This constraint makes things more difficult to optimise. One way to do that is via the method of Lagrange multipliers. To apply that, we need to reshuffle our constraint in the form of $-H_0 - \mathbb{E} \big[\ log \pi (a_t | s_t; \theta) \big] \geq 0$. Then we can form our Lagrangian at the last timestep $T$ as follow below:

Note: for this proof, we will be iterating backwards through the trajectory and finding the optimal parameters for each step. This means we have policy parameters $\theta_T$ and $\alpha_T$ that are optimal for the objective function at timestep $T$.

Assuming strong duality (i.e. that the optimal policy that solves the task has entropy $\geq H_0$), we can establish that solving our objective yields the same solution as solving our Lagrangian as:

Due to the assumed strong duality, we know that the optimal temperature would be:

Here we see a 2-step optimisation problem: first, we solve for $\pi$, then we solve for $\alpha$. Thus we can also view the above equation as:

OK, we have an expression for solving $\alpha_T$ but what about $t=0,1,\ …\ ,T-1$? Let’s go back through time using the soft-Q relationship that we previously established in Eq. 15 in order to find $\pi_{T-1}^*$:

By applying this recursively down to $t=0$ we get the generic $\alpha$ loss function:

Note: The SAC paper also features a derivation of the infinite-time version of SAC, but this blog post is concerned only with finite-time problems. That being said, I can’t understand the continuous time proof/derivation in the appendix of the paper.

Since we are now trying to solve our objective via the Lagrangian, we must constantly alternate between solving for $\theta$ and solving for $\alpha$ until convergence. However, SAC is an actor-critic method; thus, we also need to alternate between optimising for the critic parameters $\psi$. The critic here is optimised just like before. Adding that into the mix, we get to the complete SAC algorithm:

OK, so how does it perform against other algorithms? SAC outperforms DDPG across the board now. Returning to our cheetah and humanoid examples:

Cheetah

Humanoid

We can now see that SAC is not only outperforming DDPG but is also significantly more stable! On a very practical note, SAC is also easier to tune when setting up different problems. I was too lazy to leave SAC run until convergence on the humanoid task, but here are some more extensive results on all DMC tasks.

Refernces

• Reinforcement Learning: An Introduction by Sutton and Barto
• Underactuated robotics by Russ Tedrake
• Deep Deterministic Policy Gradients, Lillicrap et al. (2015)
• SAC papertrail:
  • Soft Q-Learning, Haarnoja et al. (2017)
  • Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor, Haarnoja et al. (2018)
  • Soft Actor-Critic Algorithms and Applications, Haarnoja et al. (2018)
  • Tuomas Haarnoja’s PhD thesis (2018)

Further resources

• denisyarats/pytorch_sac. A good high-performant SAC implementation.
• Lilian Weng’s blog post about policy gradients. She covers most approaches that use the theorem. Recommended!

Acknowledgements

Krishnan Srinivasan for bouncing off ideas and listening to me ramble about this

If you have any comments or suggestions, shoot me an email!