Skip to content

DDPG

lerax.algorithm.DDPG

Bases: AbstractAlgorithm[PolicyType, DDPGState[PolicyType]]

Deep Deterministic Policy Gradient (DDPG) algorithm.

An off-policy algorithm for continuous action spaces that learns a deterministic policy and a Q-function simultaneously. Uses Polyak-averaged target networks for both actor and critic.

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

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,
    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: DDPGStepState[PolicyType],
    *,
    key: Key[Array, ""],
    callback: AbstractCallback,
) -> DDPGStepState[PolicyType]

Perform a single environment step and store in replay buffer.

collect_learning_starts 

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

Collect random initial experience before training begins.

collect_rollout 

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

Collect a rollout of experience into the replay buffer.

reset 

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

iteration 

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

per_iteration 

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

Polyak-average both actor and critic target networks.

q_loss staticmethod 

q_loss(
    qf: QNetwork,
    batch: ReplayBuffer,
    target: Float[Array, " batch_size"],
) -> Float[Array, ""]

Compute MSE loss for Q-network.

actor_loss staticmethod 

actor_loss(
    policy: PolicyType, batch: ReplayBuffer, qf: QNetwork
) -> Float[Array, ""]

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

ddpg_train 

ddpg_train(
    policy: PolicyType,
    opt_state: optax.OptState,
    buffer: ReplayBuffer,
    qf: QNetwork,
    qf_target: QNetwork,
    target_policy: PolicyType,
    q_opt_state: optax.OptState,
    *,
    key: Key[Array, ""],
) -> tuple[
    PolicyType,
    optax.OptState,
    QNetwork,
    optax.OptState,
    dict[str, Scalar],
]

lerax.algorithm.DDPGState

Bases: AbstractAlgorithmState[PolicyType]

Iteration-level state for DDPG.

Attributes:

Name Type Description
iteration_count Int[Array, '']

The current iteration count.
step_state DDPGStepState[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.
qf QNetwork

The Q-network.
qf_target QNetwork

The target 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: DDPGStepState[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

qf instance-attribute 

qf: QNetwork

qf_target instance-attribute 

qf_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.