Skip to content

DQN

lerax.algorithm.DQN

Bases: AbstractAlgorithm[PolicyType, DQNState[PolicyType]]

Double Deep Q-Network (Double DQN) algorithm.

Uses the online network to select actions and the target network to evaluate them, which reduces overestimation bias compared to standard DQN.

The target network is a periodic copy of the online network, updated every target_update_interval iterations.

Attributes:

Name Type Description
optimizer optax.GradientTransformation

The optimizer used for training.
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.
target_update_interval int

How often to copy the online network to the target.
max_grad_norm float

Maximum gradient norm for gradient clipping.

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.

 100
num_envs int

Number of parallel environments.

 1
num_steps int

Number of steps per iteration.

 4
batch_size int

Batch size for training.

 32
target_update_interval int

How often to copy the online network to the target.

 10000
max_grad_norm float

Maximum gradient norm for gradient clipping.

 10.0
learning_rate optax.ScalarOrSchedule

Learning rate for the optimizer.

 0.0001

optimizer instance-attribute 

optimizer: optax.GradientTransformation = optax.chain(
    clip, adam
)

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

target_update_interval instance-attribute 

target_update_interval: int = target_update_interval

max_grad_norm instance-attribute 

max_grad_norm: float = max_grad_norm

dqn_loss_grad class-attribute instance-attribute 

dqn_loss_grad = staticmethod(
    eqx.filter_value_and_grad(dqn_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 = 1,
    num_steps: int = 4,
    batch_size: int = 32,
    target_update_interval: int = 10000,
    max_grad_norm: float = 10.0,
    learning_rate: optax.ScalarOrSchedule = 0.0001,
)

num_iterations 

num_iterations(total_timesteps: int) -> int

step 

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

Perform a single environment step and store in replay buffer.

collect_learning_starts 

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

Collect random initial experience before training begins.

collect_rollout 

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

Collect a rollout of experience into the replay buffer.

reset 

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

iteration 

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

per_iteration 

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

Periodically update the target network.

dqn_loss staticmethod 

dqn_loss(
    policy: PolicyType,
    batch: ReplayBuffer,
    target_policy: PolicyType,
    gamma: float,
) -> Float[Array, ""]

dqn_train 

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

lerax.algorithm.DQNState

Bases: AbstractAlgorithmState[PolicyType]

Iteration-level state for DQN.

Attributes:

Name Type Description
iteration_count Int[Array, '']

The current iteration count.
step_state DQNStepState[PolicyType]

The step-level state.
env AbstractEnvLike

The environment being used.
policy PolicyType

The online policy being trained.
opt_state optax.OptState

The optimizer state.
callback_state AbstractCallbackState

The callback state.
target_policy PolicyType

The target policy for stable Q-value estimation.

iteration_count instance-attribute 

iteration_count: Int[Array, '']

step_state instance-attribute 

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

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.