memorax.algorithms.DQN

Contents

memorax.algorithms.DQN#

class memorax.algorithms.DQN[source]#

Bases: object

DQN(cfg: memorax.algorithms.dqn.DQNConfig, env: gymnax.environments.environment.Environment, env_params: gymnax.environments.environment.EnvParams, q_network: flax.linen.module.Module, optimizer: optax._src.base.GradientTransformation, buffer: flashbax.buffers.trajectory_buffer.TrajectoryBuffer, epsilon_schedule: collections.abc.Callable[[typing.Union[jax.Array, numpy.ndarray, numpy.bool, numpy.number, bool, int, float, complex]], typing.Union[jax.Array, numpy.ndarray, numpy.bool, numpy.number, bool, int, float, complex]])

cfg: DQNConfig#

env: Environment#

env_params: EnvParams#

q_network: Module#

optimizer: GradientTransformation#

buffer: TrajectoryBuffer#

epsilon_schedule: Callable[[Array | ndarray | bool | number | bool | int | float | complex], Array | ndarray | bool | number | bool | int | float | complex]#

init(key)[source]#

Return type:: DQNState
Parameters:: key (Array)

warmup(key, state, num_steps)[source]#

Return type:

Parameters:

key (Array)
state (DQNState)
num_steps (int)

train(key, state, num_steps)[source]#

Return type:

Parameters:

key (Array)
state (DQNState)
num_steps (int)

evaluate(key, state, num_steps)[source]#

Return type:

Parameters:

key (Array)
state (DQNState)
num_steps (int)

__init__(cfg, env, env_params, q_network, optimizer, buffer, epsilon_schedule)#

Parameters:

cfg (DQNConfig)
env (Environment)
env_params (EnvParams)
q_network (Module)
optimizer (GradientTransformation)
buffer (TrajectoryBuffer)
epsilon_schedule (Callable[[Array | ndarray | bool | number | bool | int | float | complex], Array | ndarray | bool | number | bool | int | float | complex])

Return type:

None