Skip to content

SAC

lerax.algorithm.SAC

Bases: AbstractAlgorithm[PolicyType, SACState[PolicyType]]

Soft Actor-Critic (SAC) algorithm.

A maximum-entropy off-policy algorithm for continuous action spaces. Uses twin Q-networks with target networks, delayed policy updates, and automatic entropy coefficient tuning.

Attributes:

Name Type Description
optimizer optax.GradientTransformation

The actor optimizer.
q_optimizer optax.GradientTransformation

The Q-network optimizer.
alpha_optimizer optax.GradientTransformation

The entropy coefficient optimizer.
buffer_size int

The size of the replay buffer.
gamma float

Discount factor for future rewards.
learning_starts int

Number of initial steps to collect before training.
num_envs int

Number of parallel environments.
num_steps int

Number of steps per iteration.
batch_size int

Batch size for training.
tau float

Soft update coefficient for target networks.
policy_frequency int

How often to update the actor (relative to Q updates).
autotune bool

Whether to automatically tune the entropy coefficient.
initial_alpha float

Initial entropy coefficient value.
q_width_size int

Width of Q-network hidden layers.
q_depth int

Depth of Q-network hidden layers.

Parameters:

Name Type Description Default
buffer_size int

The size of the replay buffer.

 1000000
gamma float

Discount factor for future rewards.

 0.99
learning_starts int

Number of initial steps before training.

 100
num_envs int

Number of parallel environments.

 4
num_steps int

Number of steps per iteration.

 1
batch_size int

Batch size for training.

 256
tau float

Soft update coefficient for target networks.

 0.005
policy_frequency int

How often to update the actor.

 2
autotune bool

Whether to automatically tune entropy coefficient.

 True
initial_alpha float

Initial entropy coefficient value.

 0.2
policy_lr optax.ScalarOrSchedule

Learning rate or schedule for the actor.

 0.0003
q_lr optax.ScalarOrSchedule

Learning rate or schedule for Q-networks.

 0.0003
alpha_lr optax.ScalarOrSchedule | None

Learning rate or schedule for the entropy coefficient. Defaults to q_lr when None.

 None
q_width_size int

Width of Q-network hidden layers.

 256
q_depth int

Depth of Q-network hidden layers.

 2

optimizer instance-attribute 

optimizer: optax.GradientTransformation = (
    optax.inject_hyperparams(optax.adam)(policy_lr)
)

q_optimizer class-attribute instance-attribute 

q_optimizer: optax.GradientTransformation = (
    optax.inject_hyperparams(optax.adam)(q_lr)
)

alpha_optimizer class-attribute instance-attribute 

alpha_optimizer: optax.GradientTransformation = (
    optax.inject_hyperparams(optax.adam)(
        alpha_lr if alpha_lr is not None else q_lr
    )
)

buffer_size instance-attribute 

buffer_size: int = buffer_size

gamma instance-attribute 

gamma: float = gamma

learning_starts instance-attribute 

learning_starts: int = learning_starts

num_envs instance-attribute 

num_envs: int = num_envs

num_steps instance-attribute 

num_steps: int = num_steps

batch_size instance-attribute 

batch_size: int = batch_size

tau instance-attribute 

tau: float = tau

policy_frequency instance-attribute 

policy_frequency: int = policy_frequency

autotune instance-attribute 

autotune: bool = autotune

initial_alpha instance-attribute 

initial_alpha: float = initial_alpha

q_width_size instance-attribute 

q_width_size: int = q_width_size

q_depth instance-attribute 

q_depth: int = q_depth

q_loss_grad class-attribute instance-attribute 

q_loss_grad = staticmethod(
    eqx.filter_value_and_grad(q_loss)
)

actor_loss_grad class-attribute instance-attribute 

actor_loss_grad = staticmethod(
    eqx.filter_value_and_grad(actor_loss)
)

alpha_loss_grad class-attribute instance-attribute 

alpha_loss_grad = staticmethod(
    eqx.filter_value_and_grad(alpha_loss)
)

consolidate_callbacks 

consolidate_callbacks(
    callback: Sequence[AbstractCallback]
    | AbstractCallback
    | None = None,
) -> AbstractCallback

learn 

learn(
    env: AbstractEnvLike,
    policy: PolicyType,
    total_timesteps: int,
    *,
    key: Key[Array, ""],
    callback: Sequence[AbstractCallback]
    | AbstractCallback
    | None = None,
) -> PolicyType

Train the policy on the environment for a given number of timesteps.

Parameters:

Name Type Description Default
env AbstractEnvLike

The environment to train on.

 required
policy PolicyType

The policy to train.

 required
total_timesteps int

The total number of timesteps to train for.

 required
key Key[Array, '']

A JAX PRNG key.

 required
callback Sequence[AbstractCallback] | AbstractCallback | None

A callback or list of callbacks to use during training.

 None

Returns:

Name Type Description
policy PolicyType

The trained policy.

__init__ 

__init__(
    *,
    buffer_size: int = 1000000,
    gamma: float = 0.99,
    learning_starts: int = 100,
    num_envs: int = 4,
    num_steps: int = 1,
    batch_size: int = 256,
    tau: float = 0.005,
    policy_frequency: int = 2,
    autotune: bool = True,
    initial_alpha: float = 0.2,
    policy_lr: optax.ScalarOrSchedule = 0.0003,
    q_lr: optax.ScalarOrSchedule = 0.0003,
    alpha_lr: optax.ScalarOrSchedule | None = None,
    q_width_size: int = 256,
    q_depth: int = 2,
)

num_iterations 

num_iterations(total_timesteps: int) -> int

step 

step(
    env: AbstractEnvLike,
    policy: PolicyType,
    state: SACStepState[PolicyType],
    *,
    key: Key[Array, ""],
    callback: AbstractCallback,
) -> SACStepState[PolicyType]

Perform a single environment step and store in replay buffer.

collect_learning_starts 

collect_learning_starts(
    env: AbstractEnvLike,
    policy: PolicyType,
    step_state: SACStepState[PolicyType],
    callback: AbstractCallback,
    key: Key[Array, ""],
) -> SACStepState[PolicyType]

Collect random initial experience before training begins.

collect_rollout 

collect_rollout(
    env: AbstractEnvLike,
    policy: PolicyType,
    step_state: SACStepState[PolicyType],
    callback: AbstractCallback,
    key: Key[Array, ""],
) -> SACStepState[PolicyType]

Collect a rollout of experience into the replay buffer.

reset 

reset(
    env: AbstractEnvLike,
    policy: PolicyType,
    *,
    key: Key[Array, ""],
    callback: AbstractCallback,
) -> SACState[PolicyType]

iteration 

iteration(
    state: SACState[PolicyType],
    *,
    key: Key[Array, ""],
    callback: AbstractCallback,
) -> SACState[PolicyType]

per_iteration 

per_iteration(
    state: SACState[PolicyType],
) -> SACState[PolicyType]

Apply Polyak averaging to update target Q-networks.

q_loss staticmethod 

q_loss(
    q_params: tuple[SoftQNetwork, SoftQNetwork],
    batch: ReplayBuffer,
    target: Float[Array, " batch_size"],
) -> Float[Array, ""]

Compute combined MSE loss for twin Q-networks.

actor_loss staticmethod 

actor_loss(
    policy: PolicyType,
    batch: ReplayBuffer,
    qf1: SoftQNetwork,
    qf2: SoftQNetwork,
    alpha: Float[Array, ""],
    keys: Key[Array, " batch"],
) -> Float[Array, ""]

Compute actor loss: maximize Q-values minus entropy penalty.

alpha_loss staticmethod 

alpha_loss(
    log_alpha: Float[Array, ""],
    log_probs: Float[Array, " batch_size"],
    target_entropy: Float[Array, ""],
) -> Float[Array, ""]

Compute entropy coefficient loss.

sac_train 

sac_train(
    policy: PolicyType,
    opt_state: optax.OptState,
    buffer: ReplayBuffer,
    qf1: SoftQNetwork,
    qf2: SoftQNetwork,
    qf1_target: SoftQNetwork,
    qf2_target: SoftQNetwork,
    q_opt_state: optax.OptState,
    log_alpha: Float[Array, ""],
    alpha_opt_state: optax.OptState,
    target_entropy: Float[Array, ""],
    iteration_count: Int[Array, ""],
    *,
    key: Key[Array, ""],
) -> tuple[
    PolicyType,
    optax.OptState,
    SoftQNetwork,
    SoftQNetwork,
    optax.OptState,
    Float[Array, ""],
    optax.OptState,
    dict[str, Scalar],
]

lerax.algorithm.SACState

Bases: AbstractAlgorithmState[PolicyType]

Iteration-level state for SAC.

Attributes:

Name Type Description
iteration_count Int[Array, '']

The current iteration count.
step_state SACStepState[PolicyType]

The step-level state.
env AbstractEnvLike

The environment being used.
policy PolicyType

The policy being trained.
opt_state optax.OptState

The actor optimizer state.
callback_state AbstractCallbackState

The callback state.
qf1 SoftQNetwork

First Q-network.
qf2 SoftQNetwork

Second Q-network.
qf1_target SoftQNetwork

Target for first Q-network.
qf2_target SoftQNetwork

Target for second Q-network.
q_opt_state optax.OptState

Optimizer state for Q-networks.
log_alpha Float[Array, '']

Log of the entropy coefficient.
alpha_opt_state optax.OptState

Optimizer state for entropy coefficient.
target_entropy Float[Array, '']

Target entropy value.

iteration_count instance-attribute 

iteration_count: Int[Array, '']

step_state instance-attribute 

step_state: SACStepState[PolicyType]

env instance-attribute 

env: AbstractEnvLike

policy instance-attribute 

policy: PolicyType

opt_state instance-attribute 

opt_state: optax.OptState

callback_state instance-attribute 

callback_state: AbstractCallbackState

qf1 instance-attribute 

qf1: SoftQNetwork

qf2 instance-attribute 

qf2: SoftQNetwork

qf1_target instance-attribute 

qf1_target: SoftQNetwork

qf2_target instance-attribute 

qf2_target: SoftQNetwork

q_opt_state instance-attribute 

q_opt_state: optax.OptState

log_alpha instance-attribute 

log_alpha: Float[Array, '']

alpha_opt_state instance-attribute 

alpha_opt_state: optax.OptState

target_entropy instance-attribute 

target_entropy: Float[Array, '']

next 

next[A: AbstractAlgorithmState](
    step_state: AbstractStepState,
    policy: PolicyType,
    opt_state: optax.OptState,
) -> A

Return a new algorithm state for the next iteration.

Increments the iteration count and updates the step state, policy, and optimizer state.

Parameters:

Name Type Description Default
step_state AbstractStepState

The new step state.

 required
policy PolicyType

The new policy.

 required
opt_state optax.OptState

The new optimizer state.

 required

Returns:

Type Description
A

A new algorithm state with the updated fields.

with_callback_states 

with_callback_states[A: AbstractAlgorithmState](
    callback_state: AbstractCallbackState,
) -> A

Return a new algorithm state with the given callback state.

Parameters:

Name Type Description Default
callback_state AbstractCallbackState

The new callback state.

 required

Returns:

Type Description
A

A new algorithm state with the updated callback state.

lerax.algorithm.SoftQNetwork

Bases: eqx.Module

Soft Q-network for SAC.

Maps concatenated (observation, action) pairs to scalar Q-values using an MLP.

Attributes:

Name Type Description
mlp eqx.nn.MLP

The MLP that processes the concatenated input.

Parameters:

Name Type Description Default
observation_size int

Dimensionality of flat observations.

 required
action_size int

Dimensionality of flat actions.

 required
width_size int

Width of the hidden layers.

 256
depth int

Number of hidden layers.

 2
key Key[Array, '']

JAX PRNG key for parameter initialization.

 required

mlp instance-attribute 

mlp: eqx.nn.MLP = eqx.nn.MLP(
    in_size=observation_size + action_size,
    out_size="scalar",
    width_size=width_size,
    depth=depth,
    key=key,
)

__init__ 

__init__(
    observation_size: int,
    action_size: int,
    *,
    width_size: int = 256,
    depth: int = 2,
    key: Key[Array, ""],
)

__call__ 

__call__(
    observation: Float[Array, " obs_dim"],
    action: Float[Array, " act_dim"],
) -> Float[Array, ""]

Compute Q-value for an observation-action pair.