Agents

Every agent but TDMPC2MultiTask subclasses ActorCritic. Its constructor takes the environment, the algorithm's own hyperparameters and extensions=(); train runs one or more seeds at once. TDMPC2MultiTask learns offline from a dataset and keeps its own train, with the same seeds-at-once call.

Each agent's reference, with every constructor argument, is at the end of its algorithm page, next to the paper it follows and the notes on using it:

Agent	Class	Reference
SAC	ajax.agents.SAC.SAC.SAC	SAC
ASAC	ajax.agents.ASAC.ASAC.ASAC	ASAC
REDQ	ajax.agents.REDQ.REDQ.REDQ	REDQ
CrossQ	ajax.agents.CrossQ.CrossQ.CrossQ	CrossQ
TD3	ajax.agents.TD3.TD3.TD3	TD3
TQC	ajax.agents.TQC.TQC.TQC	TQC
IQL	ajax.agents.IQL.IQL.IQL	IQL
AVG	ajax.agents.AVG.AVG.AVG	AVG
DQN	ajax.agents.DQN.DQN.DQN	DQN
PQN	ajax.agents.PQN.PQN.PQN	PQN
PPO	ajax.agents.PPO.PPO.PPO	PPO
APO	ajax.agents.APO.APO.APO	APO
UDRL	ajax.agents.UDRL.UDRL.UDRL	UDRL
DAgger	ajax.agents.DAgger.DAgger.DAgger	DAgger
APG	ajax.agents.APG.APG.APG	APG
DreamerV3	ajax.agents.DreamerV3.DreamerV3.DreamerV3	DreamerV3
TDMPC2	ajax.agents.TDMPC2.TDMPC2.TDMPC2	TD-MPC2
TDMPC2MultiTask	ajax.agents.TDMPC2.TDMPC2MultiTask.TDMPC2MultiTask	TD-MPC2

Every agent is also exported from ajax itself (from ajax import SAC).

ajax.agents.base.ActorCritic

ActorCritic(
    env_id,
    n_envs=4,
    actor_learning_rate=0.0003,
    critic_learning_rate=0.0003,
    actor_architecture=("128", "tanh", "128", "tanh"),
    critic_architecture=("128", "tanh", "128", "tanh"),
    env_params=None,
    max_grad_norm=None,
    weight_decay=0.0,
    memory=None,
    normalize_observations=False,
    normalize_rewards=False,
    reward_normalization_gamma=None,
    actor_bias_init=None,
    critic_bias_init=None,
    encoder_kernel_init=None,
    cnn_image_shape=None,
    cnn_extra_obs_dim=0,
    cnn_spec=None,
    log_std_state_independent=False,
    log_std_init=-1.0,
    log_std_bounds=LOG_STD_BOUNDS,
    mean_kernel_init=None,
    logits_kernel_init=None,
    log_std_kernel_init=None,
    value_kernel_init=None,
    tanh_mean=False,
    encoder_layer_norm=False,
    squash=False,
    episode_length=None,
    apply_obs_normalization=True,
    action_repeat=1,
    extensions=(),
)

Initialize the agent.

Parameters:

Name	Type	Description	Default
actor_learning_rate	float | Callable[[int], float]	Learning rate for the actor's optimizer, or its schedule.	0.0003
critic_learning_rate	float | Callable[[int], float]	Learning rate for the critic's optimizer, or its schedule.	0.0003
weight_decay	float	decoupled weight decay of the actor's and the critic's optimizers (OptimizerConfig.weight_decay): on top of the Adam step, each step multiplies their parameters by 1 - lr * weight_decay. It reaches the networks only: a scalar dual variable (SAC's temperature) keeps its own plain Adam. 0 keeps plain Adam.	0.0
memory	Optional[Union[MemoryConfig, dict]]	Pluggable memory block (MemoryConfig or its dict form); None keeps the networks feedforward.	None
normalize_rewards	bool	divide the rewards the agent learns from by a running standard deviation (logged returns stay in the env's units). By default the single-step reward's.	False
reward_normalization_gamma	Optional[float]	with normalize_rewards, divide by the running standard deviation of the discounted return with this discount instead, as SB3's VecNormalize (gamma) and gymnasium's NormalizeReward do.	None
action_repeat	int	simulator steps per agent step on brax / playground envs (episode_length then counts simulator steps); stored on env_args. gymnax envs and prebuilt envs support only 1.	1

run_capabilities

run_capabilities()

The capabilities of this run (:func:~ajax.extensions.capabilities.run_capabilities).

resume_iteration_offset

resume_iteration_offset(initial_state)

Absolute index of the first scan iteration when resuming.

train feeds it to the inner train function so a resumed run sees absolute iteration indices (see build_resumable_train). The default 0 restarts the indices at every train call (every existing agent; APG's per-stage cadence relies on it). Agents whose schedules are functions of the absolute tick override this to return their tick counter -- a host-side int, equal across seeds.

get_make_train

get_make_train()

:attr:make_train with the agent's own arguments (:attr:train_kwargs) and its extensions; :meth:train passes the rest.

train

train(
    seed=42,
    n_timesteps=int(1000000.0),
    num_episode_test=10,
    logging_config=None,
    on_ids_ready=None,
    initial_state=None,
    total_timesteps=None,
    **kwargs,
)

Train the agent, every seed at once.

Parameters:

Name	Type	Description	Default
n_timesteps	int	Timesteps this call trains for.	int(1000000.0)
num_episode_test	int	Number of episodes for evaluation during training.	10
initial_state	Optional[BaseAgentState]	a state (or (state, out)) train returned, to resume from: the run continues from its timestep.	None
total_timesteps	Optional[int]	the whole run's horizon, which every schedule (training fraction, extensions, exploration, the logging's end) runs over; None: this call's end, its start (0, or the resumed timestep) plus n_timesteps. Pass the same horizon to every call of a run split into calls to follow the schedules of one call.	None

Returns (state, out), each leaf with a leading seed axis. On the shared loop (ajax.agents.loop) out is None without a logging config, else the evaluations.

train_offline

train_offline(
    dataset,
    n_updates,
    seed=42,
    num_episode_test=10,
    logging_config=None,
    on_ids_ready=None,
)

Train on a transition dataset alone, with no environment step.

n_updates of the agent's own update on a replay holding dataset's rows (:meth:ajax.agents.loop.TrainLoop.offline), evaluating every logging_config.log_frequency updates. Returns (state, out) as :meth:train does, with the state's timestep and evaluation count back at 0, so that train(initial_state=state, ...) fine-tunes online from it, its replay keeping the dataset's rows (with learning_starts=0, from the first step).

Raises ValueError on an agent without a transition replay.