model_based_rl

dreamer_v2_agent

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

Bases: OffPolicyAgent

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

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(env_fn, test_episodes: int) → 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)

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)

dreamer_v3_agent

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

Bases: OffPolicyAgent

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

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(env_fn, test_episodes: int) → 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)

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)