Teaching a Cheetah to Run with Soft Actor-Critic

A Practical Guide to Soft Actor-Critic for High-Dimensional Continuous Action Spaces

Introduction

What does it take to teach a robotic cheetah how to run? Reinforcement Learning (RL) algorithms like **Soft Actor-Critic (SAC)** give us a great start to answering this rather strange question. SAC has emerged as a standout, excelling in high-dimensional continuous action spaces with its unique blend of stability, efficiency, and exploratory capabilities.

In this guide, we’ll explore SAC and how it may be used to teach the **Half-Cheetah-v4 environment, from OpenAI Gym, to master locomotion. Along this journey, we’ll learn about hyperparameter tuning, learning curves and even get the chance to dive into real-world applications of SAC in robotics and prosthetics. Whether you’re curious about its potential or more invested in the intricacies of how it works, this guide has a little something for everyone. **Here we go!

Generated using ChatGPT-4o

Background

SAC operates within an actor-critic framework to handle continuous action spaces, leveraging two key components:

1. Critic networks: Evaluate how good an action is by estimating the Q-values (action-value function) for given state-action pairs. These networks guide the actor toward higher rewards.
2. Actor networks: Output a stochastic policy, which samples actions from a probability distribution balancing exploration and exploitation.

The off-policy nature of SAC allows it to store and reuse past experiences in a replay buffer, decoupling data collection from training. This makes SAC more sample-efficient compared to on-policy algorithms, since it doesn’t rely solely on experiences from the current policy.

Stochastic policies and Maximum entropy RL

SAC is so effective because it distinguishes itself by maintaining a stochastic policy, assigning probabilities to all possible actions rather than deterministically choosing the best one. The policy is optimised to find a balance between:

• Maximising Rewards: Ensuring the agent takes actions that lead to high cumulative rewards.
• Maximising Entropy: Encouraging the agent to explore diverse actions, reducing the risk of suboptimal strategies.

For example, in Half-Cheetah-v4, a deterministic policy might consistently apply torques in a specific pattern to the joints, missing out on more efficient ways to reach the desired goal. A stochastic policy, on the other hand, explores a wider range of joint movements and torque combinations, eventually converging on an optimal strategy for smooth and efficient locomotion.

The Maximum Entropy RL Objective augments the traditional RL reward maximisation goal with an additional entropy term:

[embed]

Where:

• π: Policy (actor network), the system that decides which actions to take.
• Q(s,a): Expected return (Q-value).
• H(π(⋅∣s)): Entropy of the policy.
• α: Temperature, which balances reward maximisation and exploration.

Soft Policy Iteration

This is the core process that controls SAC’s learning. It alternates between three main steps: policy evaluation, state-value optimisation, and policy improvement. Together, these steps help the agent understand the quality of its actions, refine its predictions, and improve its decision-making over time.

1. Policy Evaluation: Learning the Action Values

The first step involves improving the agent’s understanding of how “good” its actions are. SAC achieves this by training the Q-function, which predicts the value of taking a particular action in a given state.

This is done using a formula called the entropy-augmented Bellman equation:

[embed]

Where:

• γ is the discount factor which prioritises immediate rewards over future rewards.
• αlogπ(at+1​∣st+1​) adds a “entropy bonus”, encouraging the agent to keep exploring instead of always exploiting the highest reward.

To make the Q-function predictions more accurate, SAC minimises the difference between the predicted value and the “true” target value, where the critic networks minimise the Bellman error:

[embed]

and,

[embed]

The term Vψˉ(s′), the state-value function, represents the expected value of being in the next state s′.

2. State-Value Optimisation: Stabilising Learning

The state-value function Vψ(s) stabilises the training process by providing a baseline for comparing the Q-function’s predictions. This function estimates how good a state is, independent of the action taken.

The value function is trained to align with the predictions of the Q-function and the policy:

[embed]

This ensures that Vψ(s) accurately reflects the value of a state, factoring in both rewards and the exploration bonus (entropy).

3. Policy Improvement: Learning Better Actions

Once the agent has a better understanding of action values, it updates its policy . The goal is to maximise both the expected reward and the entropy.

The policy is trained using the following objective:

[embed]

This equation encourages the policy to favour actions with high Q-values while also maintaining enough randomness to explore alternative strategies. To compute this efficiently, SAC uses:

[embed]

Which allows gradients to flow through sampled actions, making the optimisation process smooth and stable.

4. Entropy tuning: Balancing Exploration and Exploitation

Entropy tuning dynamically adjusts the coefficient α, which determines how much randomness (entropy) is encouraged. Early in training, higher entropy helps the agent explore various strategies. As the policy improves, lower entropy allows the agent to focus on refining its actions for better performance.

The loss function for tuning α is:

[embed]

Here, Hˉ is the target entropy, which defines the desired level of exploration, it is typically equal to the negative of the number of action space dimensions.

Implementation

I implemented SAC using **Stable-Baselines3 (SB3)**, a PyTorch-based library that holds well-tested and efficient implementations of RL algorithms. SB3 simplifies the process of applying RL methods, thereby allowing us to focus on optimising SAC for the Half-Cheetah-v4 environment rather than building the algorithm from scratch.

This pseudocode narrates how SB3 implements SAC. This outlines the key steps of training, including collecting experiences, updating the critic and actor networks, dynamically tuning the entropy coefficient, and performing soft updates on target networks.

Benchmark Environment

Why Half-Cheetah-v4?

The Half-Cheetah-v4 environment presents a challenging benchmark for RL algorithms:

• High dimensionality: This environment features continuous 17-dimensional state and 6-dimensional action spaces, simulating complex robotic dynamics.
• Optimisation: The task requires learning efficient and stable running gaits for a simulated cheetah robot.
• Exploration: With diverse movement possibilities, SAC’s stochastic policy is crucial for discovering optimal strategies.

These characteristics make Half-Cheetah an ideal showcase for SAC, and its robustness and efficiency.

Half-cheetah after 1,000 steps of training

At 1,000 training steps, the Half-Cheetah wobbles rather randomly with no meaningful forward or backward motion. This illustrates early exploration, where the agent is learning basic interactions with the environment.

Half-cheetah after 100,000 steps of training

At 100,000 training steps, the Half-Cheetah has developed some semblance of control, performing an unexpected flip before it begins to ‘run’ in a particularly unconventional way. This highlights SAC’s capability to explore and evolve to understand effective movement patterns — albeit still incorrect ones.

Half-cheetah after 500,000 steps of training

By 500,000 training steps, the Half-Cheetah demonstrates refined motion. Although it starts off shaking, it quickly stabilises into a smooth running motion. This phase reflects the SAC agent’s balance between exploration and exploitation, fine-tuning its learned policy.

The shaking disappears after further training leading to near-perfect running gait.

Hyperparameter optimisation

This plays a key role in maximising the stability of SAC. Hyperparameters such as the learning rate, batch size, and entropy targets significantly influence how the algorithm balances exploration, exploitation, and stability.

Original approach

In the **original SAC paper**, the authors conducted experiments to determine ideal hyperparameters. Key parameters such as the discount factor (γ) and the target entropy were manually tuned. For instance:

• The target entropy was set based on the dimensionality of the action space −dim(action space).
• Learning rates for the policy and Q-networks were consistent across environments.
• The temperature parameter (α) was dynamically tuned to optimise the trade-off between exploration and exploitation using **Adam()**.

While effective, this approach requires significant computational resources, as hyperparameter choices often varied depending on the environment.

My Approach

To optimise SAC for the Half-Cheetah-v4 environment, I used **Optuna, a state-of-the-art hyperparameter optimisation library. It automates the process of searching for the best hyperparameters through techniques like [Bayesian optimisation](https://towardsdatascience.com/a-conceptual-explanation-of-bayesian-model-based-hyperparameter-optimization-for-machine-learning-b8172278050f)**, significantly reducing the time and effort required.

The key aspects of my implementation include:

• Learning Rate Schedule: A **cosine annealing schedule** to allow dynamic adjustment of the learning rate during training. This helps the agent to explore effectively in the early stages while fine-tuning in later stages.

• Target Entropy: I optimised the target entropy within a suitable range for the 6-dimensional action space (−7.0 to −5.0).
• Soft Update Coefficient (τ): This determines the rate at which target networks are updated for stability. A small τ ensures stable and gradual updates to the target networks, while a large τ allows faster adaptation but risks instability.
• Other Parameters: The discount factor (γ) and batch size were also tuned using Optuna to achieve a balance between stability and sample efficiency.

Results

After running 10 trials with Optuna, I found that the parameters that produced the largest cumulative reward after 100,000 steps were:

By automating hyperparameter tuning, I was able to achieve better performance in less time, demonstrating the value of systematic approaches in environments with complex dynamics like Half-Cheetah. With these parameters, the SAC model was trained for 500,000 steps — about as many steps as my hardware could handle.

Visualising the Results

This learning curve shows the SAC algorithm’s performance on Half-Cheetah-v4 over the 500,000 training steps. The agent’s mean reward rises rapidly early on, as it learns basic locomotion strategies, showcasing SAC’s effectiveness in exploration and policy improvement. By 100,000 steps, the curve smooths out, indicating a shift toward exploitation as the agent refines its running gait. In later stages, higher rewards with minimal variance highlight SAC’s stability and robustness in high-dimensional continuous control tasks.

This combined plot shows the Q-loss and policy loss during SAC training. The Q-loss decreases steadily overall, indicating the critic’s improving ability to estimate Q-values accurately, though it exhibits significant fluctuations which may be explained by the dynamic nature of training in high-dimensional action spaces. The policy loss starts near zero and becomes increasingly negative, reaching approximately -200 by the end of training. This trend reflects the actor’s gradual optimisation of actions to maximise cumulative rewards. The high variance in both losses highlights the challenges of balancing exploration and exploitation as SAC iteratively refines its policy and value estimates.

To better understand how SAC’s critic evaluates actions, I’ve laid out Q-value heatmaps for the action dimension pairs. These heatmaps reveal how the critic assigns value to combinations of actions, with higher Q-values (lighter areas) indicating optimal combinations.

For example, in the heatmap for action dimension 1 vs 2:

• High Q-Value Regions: Extreme actions (e.g. positive torque in dimension 1 and negative torque in dimension 2) are considered effective for forward motion or stabilisation — the goal.
• Low Q-Value Regions: Neutral or low-action combinations (e.g. zero torque) are less rewarding.
• Dependency: The gradients suggest that certain action combinations work synergistically to produce effective locomotion.

These heatmaps highlight the critic’s role in guiding the actor network by prioritising effective torque combinations. The smooth gradients showcase the importance of continuous action spaces for fine-tuning policies.

The action distributions generated by SAC in the Half-Cheetah environment interestingly show peaks at extreme values (±1) across all six dimensions. This pattern may reflect:

• Environment Dynamics: Extreme torque values often represent optimal strategies for rapid forward motion.
• Requirements: The task prioritises effective locomotion, where maximal torque values are often necessary to achieve stable gaits.

Despite the peaks, SAC’s stochastic policy ensures that the agent explores the action space sufficiently, allowing it to adjust its movements as training progresses.

If you’d like to test this out, my code is right here in my Github repository.

Bioengineering Context

Now that we’re here, we can safely say SAC is a versatile RL algorithm and has a wide range of applications in all aspects of life. I want to cover one in the context of bioengineering before we wrap up.

Robotic Locomotion

Just as we worked through a bipedal Half-Cheetah environment here, Haarnoja et al. demonstrates SAC’s capability in controlling a more complex quadrupedal robot, called Minitaur, enabling it to learn stable and efficient walking patterns.

Minitaur was trained and tested in environments with uneven terrains and disturbances. The learned policy allows the robot to adjust its movements dynamically, such as recovering balance after a push.

The ability to train autonomous robotic locomotion has significant implications and potential applications too. They could be utilised in a range of tasks from search-and-rescue on uneven and unstable terrains to rehabilitation or physical therapy — supporting patients with diverse needs.

On that note — we are done. Thank you for learning about SAC with me :)

I acknowledge the use of ChatGPT-4o to generate an outline for background study and display image generation; Carbon for pseudocode formatting; Embed-fun for equation formatting

I confirm that no content generated by AI has been presented as my own work.