Trainer & Agent API

Every algorithm consists of an algorithm (or trainer) and an agent. The trainer holds the hyperparameters and the training logic; the agent holds the trainable state along with the trainer that produced it.

RLAlgorithm

jaxnasium.algorithms.RLAlgorithm

Base class for reinforcement learning algorithms in JAXnasium.

This object can also be refered to as the "trainer" of an RLAgent. It contains the hyperparameters and collective methods such as train and evaluate. The trainable state lives in the RLAgent, which init_agent and train return.

The philosophy of Jaxnasium algorithms is to maintain close to single-file implementations. All training logic is therefor implemented in the files themselves, and are designed for single-agent settings. This base class provides only a default evaluation loop, environment compatibility checks, and automatic multi-agent transformation capabilities.

Key Features: - Multi-agent support: Automatic transformation of single-agent algorithms to multi-agent - Environment compatibility: Built-in checks and warnings for environment compatibility - Evaluation: Standardized evaluation interface for comparing algorithm performance

Attributes:

Name	Type	Description
agent_type	type[TAgent] | PyTree[type[TAgent]] | None	The RLAgent class that init_agent builds. In multi-agent settings, this may be a pytree with a class per agent, e.g. {"agent_0": SACAgent, "agent_1": MyAgent}, as long as the trainer supports all agent types. Optional for algorithms that override init_agent instead.
multi_agent	bool	Whether this algorithm instance operates in multi-agent mode
auto_upgrade_multi_agent	bool	Whether to automatically upgrade single-agent methods to multi-agent
log_function	Callable | Literal['simple', 'tqdm'] | None	Logging function to use ("simple", "tqdm", or custom callable)
log_interval	int | float	Interval for logging (as fraction of total steps or absolute number)
reduce_metrics_fn	Callable | Literal['mean'] | None	Reduces one training iteration's metrics to what train returns alongside the agent. This function is called after log function on the batch of iteration metrics. Defaults to "mean", averaging the metrics across the iteration. Set to None to return the full batch of metrics.

train(key: PRNGKeyArray, env: Environment, agent: Any = None) -> tuple[TAgent, PyTree[Float[Array, ' iterations']]] abstractmethod

Runs the training loop, returning the trained agent and its training metrics.

agent=None builds a fresh agent; passing an existing agent continues training it.

init_agent(key: PRNGKeyArray, env: Environment) -> TAgent

Builds and returns an agent of class agent_type for env.

evaluate(key: PRNGKeyArray, agent: RLAgent, env: Environment, num_eval_episodes: int = 10, deterministic_policy: bool = True) -> Float[Array, ' num_eval_episodes']

num_eval_episodes over the environment with the provided agent.

Can also be called through an agent via agent.evaluate(key, env, ...).

load(file_path: str) -> RLAgent staticmethod

Convenience method for RLAgent.load(file_path).

Returns the algorithm stored in file_path which was saved with save on an RLAgent, trainer included. Requires jaxon to be installed, which is not installed by default.

RLAgent

jaxnasium.algorithms.RLAgent

Trainable state, along with the trainer that produced it (hyperparameters and training logic).

trainer: eqx.AbstractVar[Any] instance-attribute

The hyperparameters this agent was built with, and which its updates read.

get_action(*args, **kwargs) -> Any abstractmethod

train(key: PRNGKeyArray, env: Environment, **hyperparams) -> tuple[Self, PyTree[Float[Array, ' iterations']]]

Continues training this agent, returning it with its training metrics.

evaluate(key: PRNGKeyArray, env: Environment, num_eval_episodes: int = 10, deterministic_policy: bool = True) -> Float[Array, ' num_eval_episodes']

with_hyperparams(**hyperparams) -> Self

Overrides this agent's trainer hyperparameters.

replace(**updates) -> Self

save(file_path: str)

Save the current state (along with the trainer) to file_path. This uses jaxon internally to save the agent. jaxon is not installed by default, so must be installed manually.

Alternatively, serialization can be done like any other eqx.Module as described here: https://docs.kidger.site/equinox/examples/serialisation/

load(file_path: str) -> Self classmethod

Returns the agent stored in file_path which was saved with save, trainer included. Requires jaxon to be installed, which is not installed by default.

Multi-agent method markers

These decorators declare how a method of an RLAgent behaves once the agent is in multi-agent mode. See Multi-Agent.

jaxnasium.algorithms._algorithm.per_agent(method)

Decorator to mark an RLAgent method as per-agent; only applicable in multi-agent settings. This is purely documentation, as it is the default behavior.

jaxnasium.algorithms._algorithm.collective(method)

Decorator to mark an RLAgent method as collective; only applicable in multi-agent settings. This means that, in multi-agent environments, this method will be called on the entire pytree of agents, rather than mapped to each agent individually.

In MA, this means that self of this method refers to the MultiAgentWrapper of all agents. This is in contrast to @per_agent, where the method is written as a single-agent method and self refers to a single RLAgent object.