Soft Actor-Critic (SAC) with JAX

This repository implements the Soft Actor-Critic (SAC) algorithm for continuous control tasks using JAX, Brax, and Equinox.

Current support:

• Only Brax environments are supported.
• Only a single instance of the HalfCheetah environment has been implemented.

The code visualizes training performance with Plotly.

Features

• JAX-based SAC implementation
• Brax environment
• Equinox neural network models
• Replay buffer
• Training and evaluation metrics
• Interactive Plotly visualization of returns

Requirements

Python 3.12+ All dependencies are managed via the uv package manager and specified in pyproject.toml.

For instructions on installing and using uv, see the official repo: https://github.com/astral-sh/uv

Install dependencies with uv:

uv sync

Usage

Run the main training script:

python sac/soft.py

Training may take a while depending on your hardware and configuration.

Hyperparameter Exploration

You can explore and experiment with different hyperparameters by modifying the config dictionary in sac/soft.py. Adjust values such as learning rates, batch size, number of epochs, and more to see their effect on training performance.

Evaluation and Plotting

After training, the script automatically evaluates the agent and generates interactive plots of average return and episode return using Plotly. Plots are shown in your browser or saved as HTML in the plots/ directory.

Current Evaluation Plot

Project Structure

value_n_policy_learning/
│
├── README.md
├── pyproject.toml
├── sac/
│   └── soft.py
└── plots/         # (created automatically for plots)

• sac/soft.py: Main SAC implementation and training loop
• pyproject.toml: Project configuration
• README.md: Project documentation
• plots/: Stores generated plots

Reference

This implementation is based on the Soft Actor-Critic algorithm. For more details, see the official paper: https://arxiv.org/abs/1801.01290

If you find this project helpful, please consider giving it a ⭐ on GitHub!