core

off_policy

class xuance.torch.agents.core.off_policy.OffPolicyAgent(config: Namespace, envs: DummyVecEnv | SubprocVecEnv | None = None, observation_space: gymnasium.spaces.Space | None = None, action_space: gymnasium.spaces.Space | None = None, callback: BaseCallback | None = None)

Bases: Agent

Base class for single-agent off-policy reinforcement learning algorithms.

This class implements the common logic shared by off-policy algorithms (e.g., DQN, DDPG, TD3, SAC) in XuanCe. It extends the generic Agent abstraction with off-policy–specific components such as replay buffers, exploration strategies, and update schedules.

The agent can be used in both training and evaluation-only scenarios. When initialized without training environments (envs=None), the agent relies on explicitly provided observation and action spaces to construct policy networks, which is useful for inference or standalone evaluation.

Parameters:

• config (Namespace) – Configuration object containing hyperparameters, algorithm settings, and runtime options.

• envs (Optional[DummyVecEnv | SubprocVecEnv]) – Vectorized environments used for training. If None, the agent will not initialize training environments and must be provided with observation_space and action_space.

• observation_space (Optional[gymnasium.spaces.Space]) – Observation space specification used to build policy and value networks when envs is None. Typically obtained from test_envs.observation_space.

• action_space (Optional[gymnasium.spaces.Space]) – Action space specification used to build policy and value networks when envs is None. Typically obtained from test_envs.action_space.

• callback (Optional[BaseCallback]) – Optional callback object for injecting custom logic during training or evaluation, such as logging, early stopping, model checkpointing, or visualization.

Notes

• Off-policy agents maintain a replay buffer to reuse past experience.

• Training and evaluation environments are conceptually separated; evaluation environments may be created and managed externally.

• In evaluation mode, exploration noise is disabled and policy actions are executed deterministically by default.

exploration(pi_actions)

Returns the actions for exploration.

Parameters:

pi_actions – The original output actions.

Returns:

The actions with noisy values.

Return type:

explore_actions

get_actions(observations: numpy.ndarray, test_mode: bool | None = False) → dict

Returns actions and values.

Parameters:

• observations (np.ndarray) – The observation.

• test_mode (Optional[bool]) – True for testing without noises.

Returns:

The actions to be executed. values: The evaluated values. dists: The policy distributions. log_pi: Log of stochastic actions.

Return type:

actions

memory: DummyOffPolicyBuffer | None

test(test_episodes: int, test_envs: DummyVecEnv | SubprocVecEnv | None = None, close_envs: bool = True) → list

Evaluate the current policy in a vectorized environment.

This method runs evaluation episodes using test_envs and returns the per-episode scores. During evaluation, actions are selected in deterministic (test) mode and optional RGB-array frames can be recorded for video logging when rendering is enabled.

Parameters:

• test_episodes (int) – Total number of evaluation episodes to run across all vectorized environments.

• test_envs (Optional[DummyVecEnv | SubprocVecEnv]) – Vectorized environments used for evaluation. Must not be None.

• close_envs (bool) – Whether to close test_envs before returning. Set this to False if test_envs is managed externally and will be reused after evaluation.

Returns:

A list of episode scores collected during evaluation.

Return type:

list

Notes

• This method resets the evaluation environments at the beginning of testing and steps them

  until test_episodes episodes are completed.

• When render_mode == “rgb_array” and self.render is True, the method records frames and logs the

  best-scoring episode as a video.

• By default, this implementation updates obs_rms during testing. If you want to avoid contaminating

  training statistics, consider guarding this update with a dedicated flag (e.g., update_rms=False).

train(train_steps: int) → dict

Run the main off-policy training loop.

This method interacts with the training environments for a fixed number of environment steps, collects transitions, stores them in the replay buffer, and periodically updates the policy using sampled mini-batches.

Training proceeds in a step-based manner (not episode-based): at each iteration, the agent selects actions, steps the environments, logs intermediate information via callbacks, and performs policy updates when the configured conditions are met.

Parameters:

train_steps (int) – Total number of environment steps to execute for training. The actual loop advances in increments of self.n_envs steps per iteration due to vectorized environments.

Returns:

A dictionary containing aggregated training information and

logged metrics collected during the training process.

Return type:

dict

Notes

• This method assumes that training environments (self.train_envs) and the replay buffer (self.memory) have already been initialized.

• Exploration behavior (e.g., epsilon-greedy or action noise) is applied during training and updated

  dynamically based on the current training step.

• Policy updates are triggered periodically according to

  self.training_frequency after self.start_training steps.

• Episode termination and reset logic are handled per environment

  instance, and episode-level statistics are logged via callbacks.

train_epochs(n_epochs=1) → dict

off_policy_marl

class xuance.torch.agents.core.off_policy_marl.OffPolicyMARLAgents(config: Namespace, envs: DummyVecMultiAgentEnv | SubprocVecMultiAgentEnv | None = None, num_agents: int | None = None, agent_keys: List[str] | None = None, state_space: gymnasium.spaces.Space | None = None, observation_space: gymnasium.spaces.Space | None = None, action_space: gymnasium.spaces.Space | None = None, callback: MultiAgentBaseCallback | None = None)

Bases: MARLAgents

Base class for multi-agent off-policy reinforcement learning algorithms.

This class implements the common logic shared by multi-agent off-policy algorithms in XuanCe. It extends the generic MARLAgents abstraction with off-policy–specific components such as replay buffers, exploration strategies (e.g., epsilon-greedy or action noise), and update schedules. It supports both feed-forward and RNN-based policies and can optionally use parameter sharing across agents.

The agent group can be used in both training and evaluation-only scenarios. When initialized without environments (envs=None), the agent group relies on explicitly provided state_space, observation_space, and action_space to build networks, which is useful for inference or standalone evaluation.

Parameters:

• config (Namespace) – Configuration object containing hyperparameters, algorithm settings, and runtime options.

• envs (Optional[DummyVecMultiAgentEnv | SubprocVecMultiAgentEnv]) – Vectorized multi-agent environments used for training. If None, the agent group will not initialize training environments and must be provided with state_space (when use_global_state=True), observation_space, and action_space.

• num_agents (Optional[int]) – Number of agents in the environment. If None, this value will be inferred from envs when available.

• agent_keys (Optional[List[str]]) – Keys/names that identify each agent in the environment. If None, inferred from envs when available.

• state_space (Optional[gymnasium.spaces.Space]) – Global state space used by centralized critics or global-state policies when enabled. Typically obtained from envs.state_space (or an equivalent field).

• observation_space (Optional[gymnasium.spaces.Space]) – Per-agent observation space specification used to construct networks when envs is None. Typically obtained from envs.observation_space.

• action_space (Optional[gymnasium.spaces.Space]) – Per-agent action space specification used to construct networks when envs is None. Typically obtained from envs.action_space.

• callback (Optional[MultiAgentBaseCallback]) – Optional callback object for injecting custom logic during training or evaluation, such as logging, early stopping, checkpointing, or visualization.

Notes

• Off-policy multi-agent agents maintain a replay buffer to reuse past experience; for RNN-based policies,

  an episode-aware buffer is used.

• Training and evaluation environments are conceptually separated; evaluation environments may be created

  and managed externally.

• In evaluation mode, exploration noise is disabled by default by setting test_mode=True when calling

  action() or run_episodes().

exploration(batch_size: int, pi_actions_dict: List[dict] | dict, avail_actions_dict: List[dict] | None = None)

Apply exploration strategy to policy actions.

This method modifies the actions produced by the policy according to the configured exploration mechanism. Supported strategies include:

• Epsilon-greedy exploration for discrete action spaces.

• Additive Gaussian noise for continuous action spaces.

The specific strategy is selected automatically based on the agent configuration (e_greedy or noise_scale).

Parameters:

• batch_size (int) – Number of parallel environments (batch size).

• pi_actions_dict (Union[List[dict], dict]) – Actions produced by the policy before exploration. When parameter sharing is enabled, this may be a shared structure across agents.

• avail_actions_dict (Optional[List[dict]]) – Available-action masks for each parallel environment when use_actions_mask=True. Can be None when action masking is disabled.

Returns:

Actions after applying the exploration strategy.

The returned structure matches the format of pi_actions_dict.

Return type:

Union[List[dict], dict]

get_actions(obs_dict: List[dict], avail_actions_dict: List[dict] | None = None, rnn_hidden: dict | None = None, test_mode: bool | None = False, **kwargs) → dict

Compute actions for all agents given vectorized observations.

This method performs a forward pass through the current multi-agent policy to obtain actions for each agent in each parallel environment. When RNN-based representations are enabled, it also consumes and returns recurrent hidden states. During training (test_mode=False), this method applies the configured exploration strategy (epsilon-greedy or additive noise); during evaluation (test_mode=True), exploration is disabled.

Parameters:

• obs_dict (List[dict]) – Observations for each parallel environment. Each element is a dict keyed by self.agent_keys.

• avail_actions_dict (Optional[List[dict]]) – Available-action masks for each parallel environment when use_actions_mask=True. Each element is a dict keyed by self.agent_keys. Can be None when action masking is disabled.

• rnn_hidden (Optional[dict]) – Current RNN hidden states keyed by self.model_keys. Required when self.use_rnn is True.

• test_mode (bool) – Whether to run in evaluation mode. When True, exploration is disabled and actions are produced deterministically (or without training-time noise).

Returns:

A dictionary containing:

• hidden_state (Optional[dict]): Updated RNN hidden states when self.use_rnn is True;

  otherwise the value returned by the policy (typically None).

• actions (List[dict]): Actions for each parallel environment.

  Each element is a dict keyed by self.agent_keys.

Return type:

dict

init_hidden_item(i_env: int, rnn_hidden: dict | None = None) → dict

Reset RNN hidden states for a specific environment index.

This method re-initializes the RNN hidden states corresponding to the i_env-th vectorized environment. When parameter sharing is enabled, the hidden state batch is arranged as (n_envs * n_agents, …), so this method resets the contiguous slice for all agents in that environment. Otherwise, it resets the single hidden-state entry for i_env for each model key.

Parameters:

• i_env (int) – Index of the vectorized environment to reset.

• rnn_hidden (Optional[dict]) – Current RNN hidden states keyed by self.model_keys. This object is updated in-place.

Returns:

Updated RNN hidden states with the i_env entries reset.

Return type:

dict

init_rnn_hidden(n_envs) → dict | None

Initialize RNN hidden states for vectorized multi-agent execution.

This method creates initial hidden states for the RNN-based policy representations

when self.use_rnn is enabled. The batch size depends on whether parameter sharing is used: - If use_parameter_sharing=True, the batch dimension is n_envs * n_agents

(one hidden state per agent per environment).

• Otherwise, the batch dimension is n_envs (one hidden state per environment per model key).

Parameters:

n_envs (int) – Number of parallel environments.

Returns:

A dictionary of initialized hidden states keyed by self.model_keys

when self.use_rnn is True; otherwise None.

Return type:

Optional[dict]

run_episodes(n_episodes: int = 1, run_envs: DummyVecMultiAgentEnv | SubprocVecMultiAgentEnv | None = None, test_mode: bool = False, close_envs: bool = True) → list

Run vectorized multi-agent episodes for rollout collection or evaluation.

This method steps a vectorized multi-agent environment using the current policy until n_episodes episodes have completed. When test_mode is False, collected transitions are stored into the replay buffer (and episode boundaries are tracked for RNN-aware buffers). When test_mode is True, exploration is disabled and episode scores are returned; optional RGB-array frames can be recorded and logged as a video.

Parameters:

• n_episodes (int) – Number of completed episodes to run across all parallel environments.

• run_envs (Optional[DummyVecMultiAgentEnv | SubprocVecMultiAgentEnv]) – Vectorized environments to run. If None, self.train_envs is used.

• test_mode (bool) – Whether to run in evaluation mode. When True, exploration is disabled and the replay buffer is not written.

• close_envs (bool) – Whether to close run_envs before returning when test_mode is True. Set this to False if the caller manages the environment lifecycle externally.

Returns:

Episode scores (mean reward across agents) for each completed episode.

Return type:

list

store_experience(obs_dict, avail_actions, actions_dict, obs_next_dict, avail_actions_next, rewards_dict, terminals_dict, info, **kwargs) → None

Store a batch of multi-agent transitions into the replay buffer.

This method converts per-environment dictionaries (one dict per vector environment) into per-agent batched arrays and writes them into the replay buffer. It also stores auxiliary fields such as agent masks and, when enabled, global state and action masks.

Parameters:

• obs_dict (List[dict]) – Observations for each parallel environment. Each element is a dict keyed by self.agent_keys.

• avail_actions (Optional[List[dict]]) – Available-action masks for each parallel environment when use_actions_mask=True. Each element is a dict keyed by self.agent_keys. Can be None when action masking is disabled.

• actions_dict (List[dict]) – Actions executed by each agent for each parallel environment. Each element is a dict keyed by self.agent_keys.

• obs_next_dict (List[dict]) – Next observations for each parallel environment. Each element is a dict keyed by self.agent_keys.

• avail_actions_next (Optional[List[dict]]) – Next-step available-action masks when use_actions_mask=True. Can be None when action masking is disabled.

• rewards_dict (List[dict]) – Rewards for each agent for each parallel environment. Each element is a dict keyed by self.agent_keys.

• terminals_dict (List[dict]) – Termination flags for each agent for each parallel environment. Each element is a dict keyed by self.agent_keys.

• info (List[dict]) – Environment info for each parallel environment at the current step. Must contain agent_mask for each agent key.

• **kwargs – Optional extra fields. When use_global_state=True, this method expects state and next_state to be provided.

test(test_episodes: int, test_envs: DummyVecMultiAgentEnv | SubprocVecMultiAgentEnv | None = None, close_envs: bool = True) → list

Evaluate the current multi-agent policy for a number of episodes.

This method runs evaluation episodes in test_envs by delegating to run_episodes(test_mode=True) and returns the per-episode scores. During evaluation, exploration is disabled and optional RGB-array frames can be recorded and logged as a video when rendering is enabled.

Parameters:

• test_episodes (int) – Number of completed episodes to evaluate across all parallel environments.

• test_envs (Optional[DummyVecMultiAgentEnv | SubprocVecMultiAgentEnv]) – Vectorized multi-agent environments used for evaluation. If None, self.train_envs is used.

• close_envs (bool) – Whether to close test_envs before returning. Set this to False if test_envs is managed externally and will be reused after evaluation.

Returns:

Episode scores (mean reward across agents) for each completed evaluation episode.

Return type:

list

train(train_steps: int) → dict

Run the main multi-agent off-policy training loop.

This method interacts with the training environments to collect multi-agent transitions, stores them in the replay buffer, and performs periodic policy updates by sampling mini-batches from the buffer. The training loop is step-based and advances in vectorized increments (one iteration corresponds to self.n_envs environment steps).

Parameters:

• train_steps (int) – Number of training iterations to run. Each iteration steps all parallel environments

• once

• self.n_envs. (so the total number of environment steps is approximately train_steps *)

Returns:

A dictionary containing aggregated training information and logged metrics collected during

training (e.g., losses, episode statistics, exploration factors).

Return type:

dict

Notes

• This method assumes that training environments (self.train_envs) and the replay buffer (self.memory)

  have already been initialized.

• When self.use_rnn is enabled, rollout collection and buffer bookkeeping are handled in run_episodes(),

  and updates are performed once enough experience is available.

• Policy updates are triggered after self.start_training steps and then periodically according to

  self.training_frequency.

train_epochs(n_epochs: int = 1) → dict

Update policies for multiple epochs using mini-batches sampled from the replay buffer.

This method performs n_epochs optimization passes. At each epoch, it samples a mini-batch from the replay buffer and calls the learner’s update function. When RNN-based policies are enabled, the RNN-specific update method is used.

Parameters:

n_epochs (int) – Number of optimization epochs to perform.

Returns:

A dictionary of training metrics returned by the learner from the last update call

(e.g., Q loss, policy loss, entropy), augmented with the current exploration factors (epsilon-greedy and noise_scale).

Return type:

dict

offline

class xuance.torch.agents.core.offline.OfflineAgent(config: Namespace, envs, callback: BaseCallback | None = None)

Bases: Agent

The core class for offline reinforcement learning.

Parameters:

• config – the Namespace variable that provides hyperparameters and other settings.

• envs – the vectorized environments.

• callback – A user-defined callback function object to inject custom logic during training. It can be used for logging, early stopping, model saving, or visualization. If not provided, a default no-op callback is used.

get_actions(observations: numpy.ndarray)

test(env_fn, steps)

train(train_steps)

train_epochs(n_epochs=1)

on_policy

class xuance.torch.agents.core.on_policy.OnPolicyAgent(config: Namespace, envs: DummyVecEnv | SubprocVecEnv | None = None, observation_space: gymnasium.spaces.Space | None = None, action_space: gymnasium.spaces.Space | None = None, callback: BaseCallback | None = None)

Bases: Agent

Base class for single-agent on-policy reinforcement learning algorithms.

This class implements the common logic shared by on-policy algorithms (e.g., A2C, PPO, TRPO) in XuanCe. It extends the generic Agent abstraction with on-policy–specific components such as trajectory buffers, rollout collection, and multi-epoch policy/value updates.

The agent can be used in both training and evaluation-only scenarios. When initialized without training environments (envs=None), the agent relies on explicitly provided observation and action spaces to construct policy networks, which is useful for inference or standalone evaluation.

Parameters:

• config (Namespace) – Configuration object containing hyperparameters, algorithm settings, and runtime options.

• envs (Optional[DummyVecEnv | SubprocVecEnv]) – Vectorized environments used for training. If None, the agent will not initialize training environments and must be provided with observation_space and action_space.

• observation_space (Optional[gymnasium.spaces.Space]) – Observation space specification used to build policy and value networks when envs is None. Typically obtained from test_envs.observation_space.

• action_space (Optional[gymnasium.spaces.Space]) – Action space specification used to build policy and value networks when envs is None. Typically obtained from test_envs.action_space.

• callback (Optional[BaseCallback]) – Optional callback object for injecting custom logic during training or evaluation, such as logging, early stopping, model checkpointing, or visualization.

Notes

• On-policy agents collect fresh trajectories from the current policy and update the policy using rollouts

  stored in a trajectory buffer.

• Training and evaluation environments are conceptually separated; evaluation environments may be created

  and managed externally.

• In evaluation mode, actions are sampled without exploration schedules specific to training

  (e.g., no epsilon-greedy / action noise).

get_actions(observations: numpy.ndarray, deterministic: bool = False, return_dists: bool = False, return_logpi: bool = False) → dict

Compute actions and value estimates for a batch of observations.

This method performs a forward pass through the current policy to obtain action distributions and value predictions. Actions are sampled stochastically from the policy distribution.

Parameters:

• observations (np.ndarray) – Batch of observations. The array is expected to have shape compatible with the underlying policy.

• deterministic (bool) – True for deterministic policy and False for stochastic policy.

• return_dists (bool) – Whether to return the action distributions (split into a Python-friendly structure).

• return_logpi (bool) – Whether to return the log-probabilities of the sampled actions.

Returns:

A dictionary containing:

• actions (np.ndarray): Sampled actions to execute in the environment(s).

• values (np.ndarray): Value estimates for the input observations.

  If the policy does not produce values, this is set to 0.

• dists (Optional[Any]): Action distributions (when return_dists=True); otherwise None.

• log_pi (Optional[np.ndarray]): Log-probabilities of sampled actions (when return_logpi=True);

  otherwise None.

Return type:

dict

get_aux_info(policy_output: dict = None) → dict

Returns auxiliary information.

Parameters:

policy_output (dict) – The output information of the policy.

Returns:

The auxiliary information.

Return type:

aux_info (dict)

get_terminated_values(observations_next: numpy.ndarray, rewards: numpy.ndarray = None) → numpy.ndarray

Compute value estimates for terminal/terminated states.

This method evaluates the value function on terminal observations and returns the value estimates used for bootstrapping (e.g., when finishing a trajectory segment).

Parameters:

• observations_next (np.ndarray) – Observations at the terminal step (or the next observations used for bootstrapping).

• rewards (Optional[np.ndarray]) – Rewards corresponding to the terminal transitions. This argument is reserved for algorithm-specific implementations and may be unused.

Returns:

Value estimates for the provided terminal observations.

Return type:

np.ndarray

test(test_episodes: int, deterministic_policy: bool = True, test_envs: DummyVecEnv | SubprocVecEnv | None = None, close_envs: bool = True) → list

Evaluate the current policy in a vectorized environment.

This method runs evaluation episodes using test_envs and returns the per-episode scores. Actions are produced by the current policy (by default sampled from the policy distribution for on-policy methods), and optional RGB-array frames can be recorded for video logging when rendering is enabled.

Parameters:

• test_episodes (int) – Total number of evaluation episodes to run across all vectorized environments.

• deterministic_policy (bool) – True for evaluating the deterministic policy, and False for evaluating the stochastic policy.

• test_envs (Optional[DummyVecEnv | SubprocVecEnv]) – Vectorized environments used for evaluation. Must not be None.

• close_envs (bool) – Whether to close test_envs before returning. Set this to False if test_envs is managed externally and will be reused after evaluation.

Returns:

A list of episode scores collected during evaluation.

Return type:

list

Notes

• This method resets the evaluation environments at the beginning of testing and steps them

  until test_episodes episodes are completed.

• When render_mode == “rgb_array” and self.render is True, the method records frames and logs

  the best-scoring episode as a video.

• By default, this implementation updates obs_rms during testing. If you want to avoid contaminating

  training statistics, consider guarding this update with a dedicated flag (e.g., update_rms=False).

train(train_steps: int) → dict

Run the main on-policy training loop.

This method interacts with the training environments to collect rollouts from the current policy, stores transitions in the on-policy trajectory buffer, and triggers policy/value updates when the buffer is full. The loop advances in vectorized steps (one iteration corresponds to self.n_envs environment steps).

Parameters:

train_steps (int) – Number of rollout collection iterations to run. Each iteration steps all vectorized environments once, so the total number of environment steps is approximately train_steps * self.n_envs.

Returns:

A dictionary containing aggregated training information and logged metrics collected during training.

Return type:

dict

Notes

• This method assumes that training environments (self.train_envs)

  and the trajectory buffer (self.memory) have already been initialized.

• After collecting horizon_size steps per environment, the buffer becomes full and the agent computes

  bootstrapped terminal values, finalizes trajectory segments via finish_path, and performs n_epochs optimization passes over mini-batches using train_epochs.

• Episode termination and reset logic are handled per environment,

  and episode-level statistics are reported via callbacks.

train_epochs(n_epochs: int = 1) → dict

Update the policy for multiple epochs using samples from the rollout buffer.

This method performs multiple passes over the collected rollout data in self.memory. For each epoch, it shuffles transition indices and iterates over mini-batches to compute gradient updates via the learner.

Parameters:

n_epochs (int) – Number of optimization epochs to perform over the current rollout buffer.

Returns:

A dictionary of training metrics returned by the learner from the last mini-batch update

(e.g., policy loss, value loss, entropy, KL divergence). Implementations may include additional diagnostics depending on the algorithm.

Return type:

dict

on_policy_marl

class xuance.torch.agents.core.on_policy_marl.OnPolicyMARLAgents(config: Namespace, envs: DummyVecMultiAgentEnv | SubprocVecMultiAgentEnv | None = None, num_agents: int | None = None, agent_keys: List[str] | None = None, state_space: gymnasium.spaces.Space | None = None, observation_space: gymnasium.spaces.Space | None = None, action_space: gymnasium.spaces.Space | None = None, callback: MultiAgentBaseCallback | None = None)

Bases: MARLAgents

Base class for multi-agent on-policy reinforcement learning algorithms.

This class implements the common logic shared by multi-agent on-policy algorithms (e.g., MAPPO, MATrPO, and other actor-critic variants) in XuanCe. It extends the generic MARLAgents abstraction with on-policy–specific components such as trajectory buffers, rollout collection, advantage/return estimation (GAE), and multi-epoch policy/value updates.

The agent group can be used in both training and evaluation-only scenarios. When initialized without environments (envs=None), the agent group relies on explicitly provided state_space, observation_space, and action_space to build networks, which is useful for inference or standalone evaluation.

Parameters:

• config (Namespace) – Configuration object containing hyperparameters, algorithm settings, and runtime options.

• envs (Optional[DummyVecMultiAgentEnv | SubprocVecMultiAgentEnv]) – Vectorized multi-agent environments used for training. If None, the agent group will not initialize training environments and must be provided with state_space (when use_global_state=True), observation_space, and action_space.

• num_agents (Optional[int]) – Number of agents in the environment. If None, this value will be inferred from envs when available.

• agent_keys (Optional[List[str]]) – Keys/names that identify each agent in the environment. If None, inferred from envs when available.

• state_space (Optional[gymnasium.spaces.Space]) – Global state space used by centralized critics or global-state policies when enabled. Typically obtained from envs.state_space (or an equivalent field).

• observation_space (Optional[gymnasium.spaces.Space]) – Per-agent observation space specification used to construct networks when envs is None. Typically obtained from envs.observation_space.

• action_space (Optional[gymnasium.spaces.Space]) – Per-agent action space specification used to construct networks when envs is None. Typically obtained from envs.action_space.

• callback (Optional[MultiAgentBaseCallback]) – Optional callback object for injecting custom logic during training or evaluation, such as logging, early stopping, checkpointing, or visualization.

Notes

• On-policy multi-agent agents collect fresh rollouts from the current policy and update the policy

  using trajectories stored in a buffer.

• Training and evaluation environments are conceptually separated; evaluation environments may be created and

  managed externally.

• In evaluation mode, exploration schedules specific to training are disabled by default

  (e.g., actions are sampled without epsilon-greedy or additive noise used by off-policy agents).

get_actions(obs_dict: List[dict], state: numpy.ndarray | None = None, avail_actions_dict: List[dict] | None = None, rnn_hidden_actor: dict | None = None, rnn_hidden_critic: dict | None = None, test_mode: bool | None = False, deterministic: bool = False, **kwargs) → dict

Compute actions (and optional value/log-prob outputs) for multi-agent execution.

This method performs a forward pass through the current multi-agent actor-critic policy to produce actions for each agent in each parallel environment. When RNN-based representations are enabled, the method consumes and returns recurrent hidden states for both the actor and the critic. During training (test_mode=False), this method also computes critic values and action log-probabilities needed for on-policy updates. During evaluation (test_mode=True), critic values and log-probabilities are not computed to reduce overhead.

Parameters:

• obs_dict (List[dict]) – Observations for each parallel environment. Each element is a dict keyed by self.agent_keys.

• state (Optional[np.ndarray]) – Global state array used by centralized critics when use_global_state=True. The expected shape depends on the environment wrapper.

• avail_actions_dict (Optional[List[dict]]) – Available-action masks for each parallel environment when use_actions_mask=True. Each element is a dict keyed by self.agent_keys. Can be None when action masking is disabled.

• rnn_hidden_actor (Optional[dict]) – Current actor RNN hidden states keyed by self.model_keys. Required when self.use_rnn is True.

• rnn_hidden_critic (Optional[dict]) – Current critic RNN hidden states keyed by self.model_keys. Required when self.use_rnn is True and values are requested.

• test_mode (bool) – Whether to run in evaluation mode. When True, only actions are produced and training-specific outputs (values/log_pi) are omitted.

• deterministic (bool) – True for deterministic policy and False for stochastic policy.

Returns:

A dictionary containing:

• rnn_hidden_actor (Optional[dict]): Updated actor RNN hidden states when self.use_rnn is True;

  otherwise the value returned by the policy (typically None).

• rnn_hidden_critic (Optional[dict]): Updated critic RNN hidden states when computed;

  otherwise an empty dict.

• actions (List[dict]): Actions for each parallel environment. Each element is a dict keyed by

  self.agent_keys.

• log_pi (dict): Log-probabilities of sampled actions for each agent when test_mode=False;

  otherwise an empty dict.

• values (dict): Critic value estimates for each agent when test_mode=False; otherwise an empty dict.

Return type:

dict

init_hidden_item(i_env: int, rnn_hidden_actor: dict | None = None, rnn_hidden_critic: dict | None = None) → Tuple[dict | None, dict | None]

Reset RNN hidden states for a specific environment index.

This method re-initializes the RNN hidden states corresponding to the i_env-th vectorized environment. When parameter sharing is enabled, the hidden state batch is arranged as (n_envs * n_agents, …), so this method resets the contiguous slice for all agents in that environment. Otherwise, it resets the single hidden-state entry for i_env for each model key.

Parameters:

• i_env (int) – Index of the vectorized environment to reset.

• rnn_hidden_actor (Optional[dict]) – Current actor RNN hidden states keyed by self.model_keys. This object is updated in-place.

• rnn_hidden_critic (Optional[dict]) – Current critic RNN hidden states keyed by self.model_keys. This object is updated in-place. Can be None when critic hidden states are not tracked.

Returns:

Updated (rnn_hidden_actor, rnn_hidden_critic) with

the i_env entries reset.

Return type:

Tuple[Optional[dict], Optional[dict]]

init_rnn_hidden(n_envs) → Tuple[dict | None, dict | None]

Initialize RNN hidden states for vectorized multi-agent execution.

This method creates initial hidden states for the RNN-based actor and critic representations when self.use_rnn is enabled. The batch size depends on whether parameter sharing is used:

• If use_parameter_sharing=True, the batch dimension is n_envs * n_agents

  (one hidden state per agent per environment).

• Otherwise, the batch dimension is n_envs (one hidden state per environment per model key).

Parameters:

n_envs (int) – Number of parallel environments.

Returns:

A tuple of (rnn_hidden_actor, rnn_hidden_critic).

Each element is a dict keyed by self.model_keys when self.use_rnn is True; otherwise both are None.

Return type:

Tuple[Optional[dict], Optional[dict]]

run_episodes(n_episodes: int = 1, deterministic_policy: bool = False, run_envs: DummyVecMultiAgentEnv | SubprocVecMultiAgentEnv | None = None, test_mode: bool = False, close_envs: bool = True) → list

Run vectorized multi-agent episodes for rollout collection or evaluation.

This method steps a vectorized multi-agent environment using the current actor-critic policy until n_episodes episodes have completed. When test_mode is False, collected transitions are stored into the on-policy trajectory buffer and episode boundaries are tracked for bootstrapping and advantage computation (GAE). When test_mode is True, training-time outputs (values/log-probabilities) are skipped, exploration schedules are disabled by default, and episode scores are returned; optional RGB-array frames can be recorded and logged as a video.

Parameters:

• n_episodes (int) – Number of completed episodes to run across all parallel environments.

• deterministic_policy (bool) – True for evaluating the deterministic policy, and False for evaluating the stochastic policy.

• run_envs (Optional[DummyVecMultiAgentEnv | SubprocVecMultiAgentEnv]) – Vectorized environments to run. If None, self.train_envs is used.

• test_mode (bool) – Whether to run in evaluation mode. When True, the trajectory buffer is not written and only episode scores are collected.

• close_envs (bool) – Whether to close run_envs before returning when test_mode is True. Set this to False if the caller manages the environment lifecycle externally.

Returns:

Episode scores (mean reward across agents) for each completed episode.

Return type:

list

store_experience(obs_dict, avail_actions, actions_dict, log_pi_a, rewards_dict, values_dict, terminals_dict, info, **kwargs) → None

Store a batch of multi-agent transitions into the on-policy buffer.

This method converts per-environment dictionaries (one dict per vector environment) into per-agent batched arrays and writes them into the on-policy trajectory buffer. It also stores auxiliary fields such as agent masks and (optionally) global state and action masks. For RNN-based policies, episode-step indices are recorded to support episode-aware bookkeeping.

Parameters:

• obs_dict (List[dict]) – Observations for each parallel environment. Each element is a dict keyed by self.agent_keys.

• avail_actions (Optional[List[dict]]) – Available-action masks for each parallel environment when use_actions_mask=True. Each element is a dict keyed by self.agent_keys. Can be None when action masking is disabled.

• actions_dict (List[dict]) – Actions executed by each agent for each parallel environment. Each element is a dict keyed by self.agent_keys.

• log_pi_a (dict) – Log-probabilities of the actions under the current policy (typically computed during rollout collection).

• rewards_dict (List[dict]) – Rewards for each agent for each parallel environment. Each element is a dict keyed by self.agent_keys.

• values_dict (dict) – Value estimates produced by the critic for each agent (used for advantage/return computation).

• terminals_dict (List[dict]) – Termination flags for each agent for each parallel environment. Each element is a dict keyed by self.agent_keys.

• info (List[dict]) – Environment info for each parallel environment at the current step. Must contain agent_mask for each agent key.

• **kwargs – Optional extra fields. When use_global_state=True, this method expects state to be provided.

test(test_episodes: int, deterministic_policy: bool = True, test_envs: DummyVecMultiAgentEnv | SubprocVecMultiAgentEnv | None = None, close_envs: bool = True) → list

Evaluate the current multi-agent policy for a number of episodes.

This method runs evaluation episodes in test_envs by delegating to run_episodes(test_mode=True) and returns the per-episode scores. During evaluation, training-time outputs are skipped, exploration schedules are disabled by default, and optional RGB-array frames can be recorded and logged as a video when rendering is enabled.

Parameters:

• test_episodes (int) – Number of completed episodes to evaluate across all parallel environments.

• deterministic_policy (bool) – True for evaluating the deterministic policy, and False for evaluating the stochastic policy.

• test_envs (Optional[DummyVecMultiAgentEnv | SubprocVecMultiAgentEnv]) – Vectorized multi-agent environments used for evaluation. If None, self.train_envs is used.

• close_envs (bool) – Whether to close test_envs before returning. Set this to False if test_envs is managed externally and will be reused after evaluation.

Returns:

Episode scores (mean reward across agents) for each completed evaluation episode.

Return type:

list

train(train_steps: int) → dict

Run the main multi-agent on-policy training loop.

This method interacts with the training environments to collect fresh rollouts from the current policy, stores transitions in the on-policy trajectory buffer, and triggers policy/value updates when the buffer is full. Training advances in vectorized increments (one iteration corresponds to stepping all parallel environments once).

Parameters:

train_steps (int) – Number of rollout collection iterations to run. Each iteration steps all parallel environments once, so the total number of environment steps is approximately train_steps * self.n_envs.

Returns:

A dictionary containing aggregated training information and logged metrics collected during

training (e.g., policy loss, value loss, entropy, KL divergence, and episode statistics).

Return type:

dict

Notes

• This method assumes that training environments (self.train_envs) and the trajectory buffer self.memory

  have already been initialized.

• When the buffer becomes full, the agent finalizes trajectories by computing bootstrapped terminal values

  via values_next and calling finish_path, then performs n_epochs optimization passes over mini-batches using train_epochs.

• Episode termination and reset logic are handled per environment,

  and episode-level statistics are reported via callbacks.

train_epochs(n_epochs: int = 1) → dict

Update policies for multiple epochs using mini-batches from the trajectory buffer.

This method performs n_epochs optimization passes over the rollout data stored in self.memory. For each epoch, it shuffles transition indices and iterates over mini-batches to compute gradient updates via the learner. When RNN-based policies are enabled, the RNN-specific update method is used.

Parameters:

n_epochs (int) – Number of optimization epochs to perform over the current trajectory buffer.

Returns:

A dictionary of training metrics returned by the learner from the last mini-batch update (e.g., policy

loss, value loss, entropy, KL divergence). Implementations may include additional diagnostics depending on the algorithm.

Return type:

dict

values_next(i_env: int, obs_dict: dict, state: numpy.ndarray | None = None, rnn_hidden_critic: dict | None = None) → Tuple[dict | None, dict | None]

Compute bootstrapped critic values for an environment that reached a boundary.

This method evaluates the critic on the terminal/next observations of a specific vectorized environment (i_env) and returns per-agent value estimates used for bootstrapping when finalizing trajectories (e.g., for GAE/return computation).

Parameters:

• i_env (int) – Index of the vectorized environment that is finishing an episode or trajectory segment.

• obs_dict (dict) – Per-agent observations for the selected environment. This dict is keyed by self.agent_keys.

• state (Optional[np.ndarray]) – Global state for the selected environment when use_global_state=True. If provided, it should correspond to the same i_env instance.

• rnn_hidden_critic (Optional[dict]) – Current critic RNN hidden states keyed by self.model_keys. Required when self.use_rnn is True.

Returns:

A tuple of (rnn_hidden_critic_new, values_dict):

• rnn_hidden_critic_new (Optional[dict]): Updated critic hidden states for the selected environment

  when self.use_rnn is True; otherwise the value returned by the critic (typically None).

• values_dict (dict): Per-agent critic value estimates keyed by self.agent_keys.

Return type:

Tuple[Optional[dict], dict]