Skip to content

TD3

lerax.algorithm.TD3

Bases: AbstractAlgorithm[PolicyType, TD3State[PolicyType]]

Twin Delayed DDPG (TD3) algorithm.

Extends DDPG with three improvements: twin Q-networks (take the minimum for target computation to reduce overestimation), delayed policy updates (actor updates less frequently than critic), and target policy smoothing (clipped noise added to target actions).

Attributes:

Name Type Description
optimizer optax.GradientTransformation

The actor optimizer.
q_optimizer optax.GradientTransformation

The Q-network 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).
policy_noise float

Noise scale for target policy smoothing.
noise_clip float

Clipping range for target policy smoothing noise.
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 to collect before training.

 25000
num_envs int

Number of parallel environments.

 1
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
policy_noise float

Noise scale for target policy smoothing.

 0.2
noise_clip float

Clipping range for target policy smoothing noise.

 0.5
actor_lr optax.ScalarOrSchedule

Learning rate for the actor.

 0.0003
q_lr optax.ScalarOrSchedule

Learning rate for Q-networks.

 0.0003
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)(actor_lr)
)

q_optimizer class-attribute instance-attribute 

q_optimizer: optax.GradientTransformation = (
    optax.inject_hyperparams(optax.adam)(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

policy_noise instance-attribute 

policy_noise: float = policy_noise

noise_clip instance-attribute 

noise_clip: float = noise_clip

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)
)

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 = 25000,
    num_envs: int = 1,
    num_steps: int = 1,
    batch_size: int = 256,
    tau: float = 0.005,
    policy_frequency: int = 2,
    policy_noise: float = 0.2,
    noise_clip: float = 0.5,
    actor_lr: optax.ScalarOrSchedule = 0.0003,
    q_lr: optax.ScalarOrSchedule = 0.0003,
    q_width_size: int = 256,
    q_depth: int = 2,
)

num_iterations 

num_iterations(total_timesteps: int) -> int

step 

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

Perform a single environment step and store in replay buffer.

collect_learning_starts 

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

Collect random initial experience before training begins.

collect_rollout 

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

Collect a rollout of experience into the replay buffer.

reset 

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

iteration 

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

per_iteration 

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

Polyak-average actor and critic target networks.

q_loss staticmethod 

q_loss(
    q_params: tuple[QNetwork, QNetwork],
    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: QNetwork
) -> Float[Array, ""]

Compute actor loss: maximize Q1-value of deterministic action.

td3_train 

td3_train(
    policy: PolicyType,
    opt_state: optax.OptState,
    buffer: ReplayBuffer,
    qf1: QNetwork,
    qf2: QNetwork,
    qf1_target: QNetwork,
    qf2_target: QNetwork,
    target_policy: PolicyType,
    q_opt_state: optax.OptState,
    iteration_count: Int[Array, ""],
    *,
    key: Key[Array, ""],
) -> tuple[
    PolicyType,
    optax.OptState,
    QNetwork,
    QNetwork,
    optax.OptState,
    dict[str, Scalar],
]

lerax.algorithm.TD3State

Bases: AbstractAlgorithmState[PolicyType]

Iteration-level state for TD3.

Attributes:

Name Type Description
iteration_count Int[Array, '']

The current iteration count.
step_state TD3StepState[PolicyType]

The step-level state.
env AbstractEnvLike

The environment being used.
policy PolicyType

The online policy being trained.
opt_state optax.OptState

The actor optimizer state.
callback_state AbstractCallbackState

The callback state.
target_policy PolicyType

The target actor for stable Q-value estimation.
qf1 QNetwork

First Q-network.
qf2 QNetwork

Second Q-network.
qf1_target QNetwork

Target for first Q-network.
qf2_target QNetwork

Target for second Q-network.
q_opt_state optax.OptState

The Q-network optimizer state.

iteration_count instance-attribute 

iteration_count: Int[Array, '']

step_state instance-attribute 

step_state: TD3StepState[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

target_policy instance-attribute 

target_policy: PolicyType

qf1 instance-attribute 

qf1: QNetwork

qf2 instance-attribute 

qf2: QNetwork

qf1_target instance-attribute 

qf1_target: QNetwork

qf2_target instance-attribute 

qf2_target: QNetwork

q_opt_state instance-attribute 

q_opt_state: optax.OptState

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.