model_based_rl

dreamer_v2_agent

class xuance.torch.agents.model_based_rl.dreamer_v2_agent.DreamerV2Agent(config: Namespace, envs: DummyVecEnv | SubprocVecEnv, callback: BaseCallback | None = None)[source]

Bases: OffPolicyAgent

get_actions(obs: numpy.ndarray, test_mode: bool | None = False, player: PlayerDV2 | None = None) numpy.ndarray[source]

Returns actions and values.

Parameters:
  • obs (np.ndarray) – The observation.

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

  • player (Optional[PlayerDV2]) – The player whose action is taken, default is train_player.

Returns:

The real_actions to be executed.

Return type:

actions

obs_shape

hwc 2 chw: agent & memory both uses ‘hwc’ obs needed to be transformed to ‘chw’ and be normalized before sample & taking an action

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

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)[source]

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: int = 1)[source]

dreamer_v3_agent

class xuance.torch.agents.model_based_rl.dreamer_v3_agent.DreamerV3Agent(config: Namespace, envs: DummyVecEnv | SubprocVecEnv, callback: BaseCallback | None = None)[source]

Bases: OffPolicyAgent

get_actions(obs: numpy.ndarray, test_mode: bool | None = False, player: PlayerDV3 | None = None) numpy.ndarray[source]

Returns actions and values.

Parameters:
  • obs (np.ndarray) – The observation.

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

  • player (Optional[PlayerDV3]) – The player whose action is taken, default is train_player.

Returns:

The real_actions to be executed.

Return type:

actions

obs_shape

hwc 2 chw: agent & memory both uses ‘hwc’ obs needed to be transformed to ‘chw’ and be normalized before sample & taking an action

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

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)[source]

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: int = 1)[source]