Stable Baselines3 Contrib

https://sb3-contrib.readthedocs.io/en/master/
Admin
Stable Baselines3 Contrib is an experimental contribution package for Stable Baselines3 that...
Tokens:231,791
Snippets:1,048
Trust Score:9.6
Update:2 days ago
Show doc for...
Tokens:
Raw
### Install MicroRTS and dependencies

Source: https://sb3-contrib.readthedocs.io/en/master/modules/ppo_mask.html

Install MicroRTS by downloading and unzipping a specific artifact, then install Python requirements. Optionally create a virtual environment before installing packages.

```bash
# Install MicroRTS:
rm -fR ~/microrts && mkdir ~/microrts && \
 wget -O ~/microrts/microrts.zip http://microrts.s3.amazonaws.com/microrts/artifacts/202004222224.microrts.zip && \
 unzip ~/microrts/microrts.zip -d ~/microrts/

# You may want to make a venv before installing packages
pip install -r requirements.txt
```

--------------------------------

### Install bleeding-edge version from GitHub

Source: https://sb3-contrib.readthedocs.io/en/master/guide/install.html

Installs the latest development version directly from the GitHub repository.

```bash
pip install git+https://github.com/Stable-Baselines-Team/stable-baselines3-contrib/
```

--------------------------------

### Clone and install development version

Source: https://sb3-contrib.readthedocs.io/en/master/guide/install.html

Clones the repository and installs it in editable mode for development with test and documentation support.

```bash
git clone https://github.com/Stable-Baselines-Team/stable-baselines3-contrib/ && cd stable-baselines3-contrib
pip install -e .
```

--------------------------------

### Install stable release with pip

Source: https://sb3-contrib.readthedocs.io/en/master/guide/install.html

Installs the stable release of Stable Baselines3 Contrib using pip.

```bash
pip install sb3-contrib
```

--------------------------------

### learn method implementation

Source: https://sb3-contrib.readthedocs.io/en/master/_modules/stable_baselines3/common/on_policy_algorithm.html

Main learning loop that sets up learning, calls the training start callback, and repeatedly collects rollouts until the total timesteps are reached or training is interrupted.

```python
    def learn(
        self: SelfOnPolicyAlgorithm,
        total_timesteps: int,
        callback: MaybeCallback = None,
        log_interval: int = 1,
        tb_log_name: str = "OnPolicyAlgorithm",
        reset_num_timesteps: bool = True,
        progress_bar: bool = False,
    ) -> SelfOnPolicyAlgorithm:
        iteration = 0

        total_timesteps, callback = self._setup_learn(
            total_timesteps,
            callback,
            reset_num_timesteps,
            tb_log_name,
            progress_bar,
        )

        callback.on_training_start(locals(), globals())

        assert self.env is not None

        while self.num_timesteps < total_timesteps:
            continue_training = self.collect_rollouts(self.env, callback, self.rollout_buffer, n_rollout_steps=self.n_steps)

            if not continue_training:
                break
```

--------------------------------

### CrossQ usage example

Source: https://sb3-contrib.readthedocs.io/en/master/modules/crossq.html

Shows how to create, train, and save a CrossQ model with the MlpPolicy on the Walker2d-v4 environment.

```python
from sb3_contrib import CrossQ

model = CrossQ("MlpPolicy", "Walker2d-v4")
model.learn(total_timesteps=1_000_000)
model.save("crossq_walker")
```

--------------------------------

### Set up learning with _setup_learn

Source: https://sb3-contrib.readthedocs.io/en/master/_modules/stable_baselines3/common/base_class.html

Initializes variables needed for training: sets start time, resets buffers and counters if needed, resets action noise, handles total timesteps, resets environment, configures logger, and initializes callbacks. Returns total timesteps and callback.

```python
self.start_time = time.time_ns()

if self.ep_info_buffer is None or reset_num_timesteps:
    # Initialize buffers if they don't exist, or reinitialize if resetting counters
    self.ep_info_buffer = deque(maxlen=self._stats_window_size)
    self.ep_success_buffer = deque(maxlen=self._stats_window_size)

if self.action_noise is not None:
    self.action_noise.reset()

if reset_num_timesteps:
    self.num_timesteps = 0
    self._episode_num = 0
else:
    # Make sure training timesteps are ahead of the internal counter
    total_timesteps += self.num_timesteps
self._total_timesteps = total_timesteps
self._num_timesteps_at_start = self.num_timesteps

# Avoid resetting the environment when calling ``.learn()`` consecutive times
if reset_num_timesteps or self._last_obs is None:
    assert self.env is not None
    self._last_obs = self.env.reset()  # type: ignore[assignment]
    self._last_episode_starts = np.ones((self.env.num_envs,), dtype=bool)
    # Retrieve unnormalized observation for saving into the buffer
    if self._vec_normalize_env is not None:
        self._last_original_obs = self._vec_normalize_env.get_original_obs()

# Configure logger's outputs if no logger was passed
if not self._custom_logger:
    self._logger = utils.configure_logger(self.verbose, self.tensorboard_log, tb_log_name, reset_num_timesteps)

# Create eval callback if needed
callback = self._init_callback(callback, progress_bar)

return total_timesteps, callback
```

--------------------------------

### QRDQN _setup_model method

Source: https://sb3-contrib.readthedocs.io/en/master/_modules/sb3_contrib/qrdqn/qrdqn.html

Sets up the model by calling the parent setup, creating aliases, copying running stats for batch normalization, and initializing the exploration schedule. Warns if the number of environments exceeds the target update interval, which would cause more frequent target updates.

```python
    def _setup_model(self) -> None:
        super()._setup_model()
        self._create_aliases()
        # Copy running stats, see https://github.com/DLR-RM/stable-baselines3/issues/996
        self.batch_norm_stats = get_parameters_by_name(self.quantile_net, ["running_"])
        self.batch_norm_stats_target = get_parameters_by_name(self.quantile_net_target, ["running_"])
        self.exploration_schedule = LinearSchedule(
            self.exploration_initial_eps, self.exploration_final_eps, self.exploration_fraction
        )

        if self.n_envs > 1:
            if self.n_envs > self.target_update_interval:
                warnings.warn(
                    "The number of environments used is greater than the target network "
                    f"update interval ({self.n_envs} > {self.target_update_interval}), "
                    "therefore the target network will be updated after each call to env.step() "
                    f"which corresponds to {self.n_envs} steps."
                )
```

--------------------------------

### CrossQ _setup_model method

Source: https://sb3-contrib.readthedocs.io/en/master/_modules/sb3_contrib/crossq/crossq.html

Sets up the model, creates aliases, and configures target entropy and entropy coefficient. Handles automatic target entropy and learned entropy coefficient.

```python
    def _setup_model(self) -> None:
        super()._setup_model()
        self._create_aliases()
        # Target entropy is used when learning the entropy coefficient
        if self.target_entropy == "auto":
            # automatically set target entropy if needed
            self.target_entropy = float(-np.prod(self.env.action_space.shape).astype(np.float32))  # type: ignore
        else:
            # Force conversion
            # this will also throw an error for unexpected string
            self.target_entropy = float(self.target_entropy)

        # The entropy coefficient or entropy can be learned automatically
        # see Automating Entropy Adjustment for Maximum Entropy RL section
        # of https://arxiv.org/abs/1812.05905
        if isinstance(self.ent_coef, str) and self.ent_coef.startswith("auto"):
            # Default initial value of ent_coef when learned
            init_value = 1.0
            if "_" in self.ent_coef:
                init_value = float(self.ent_coef.split("_")[1])
                assert init_value > 0.0, "The initial value of ent_coef must be greater than 0"

            # Note: we optimize the log of the entropy coeff which is slightly different from the paper
            # as discussed in https://github.com/rail-berkeley/softlearning/issues/37
            self.log_ent_coef = th.log(th.ones(1, device=self.device) * init_value).requires_grad_(True)
            self.ent_coef_optimizer = th.optim.Adam([self.log_ent_coef], lr=self.lr_schedule(1))
        else:
            # Force conversion to float
            # this will throw an error if a malformed string (different from 'auto')
            # is passed
            self.ent_coef_tensor = th.tensor(float(self.ent_coef), device=self.device)
```

--------------------------------

### OnPolicyAlgorithm _setup_model method

Source: https://sb3-contrib.readthedocs.io/en/master/_modules/stable_baselines3/common/on_policy_algorithm.html

Sets up learning rate schedule, random seed, rollout buffer (choosing DictRolloutBuffer for Dict observation spaces), and policy. Moves policy to device and warns about CPU usage.

```python
    def _setup_model(self) -> None:
        self._setup_lr_schedule()
        self.set_random_seed(self.seed)

        if self.rollout_buffer_class is None:
            if isinstance(self.observation_space, spaces.Dict):
                self.rollout_buffer_class = DictRolloutBuffer
            else:
                self.rollout_buffer_class = RolloutBuffer

        self.rollout_buffer = self.rollout_buffer_class(
            self.n_steps,
            self.observation_space,  # type: ignore[arg-type]
            self.action_space,
            device=self.device,
            gamma=self.gamma,
            gae_lambda=self.gae_lambda,
            n_envs=self.n_envs,
            **self.rollout_buffer_kwargs,
        )
        self.policy = self.policy_class(  # type: ignore[assignment]
            self.observation_space, self.action_space, self.lr_schedule, use_sde=self.use_sde, **self.policy_kwargs
        )
        self.policy = self.policy.to(self.device)
        # Warn when not using CPU with MlpPolicy
        self._maybe_recommend_cpu()
```

--------------------------------

### Process sequence through LSTM with episode starts

Source: https://sb3-contrib.readthedocs.io/en/master/_modules/sb3_contrib/common/recurrent/policies.html

Static method that performs a forward pass in the LSTM network, handling episode starts by resetting states when necessary. It reshapes features and episode starts, and uses a fast path when no resets are needed.

```python
    @staticmethod
    def _process_sequence(
        features: th.Tensor,
        lstm_states: tuple[th.Tensor, th.Tensor],
        episode_starts: th.Tensor,
        lstm: nn.LSTM,
    ) -> tuple[th.Tensor, th.Tensor]:
        """
        Do a forward pass in the LSTM network.

        :param features: Input tensor
        :param lstm_states: previous hidden and cell states of the LSTM, respectively
        :param episode_starts: Indicates when a new episode starts,
            in that case, we need to reset LSTM states.
        :param lstm: LSTM object.
        :return: LSTM output and updated LSTM states.
        """
        # LSTM logic
        # (sequence length, batch size, features dim)
        # (batch size = n_envs for data collection or n_seq when doing gradient update)
        n_seq = lstm_states[0].shape[1]
        # Batch to sequence
        # (padded batch size, features_dim) -> (n_seq, max length, features_dim) -> (max length, n_seq, features_dim)
        # note: max length (max sequence length) is always 1 during data collection
        features_sequence = features.reshape((n_seq, -1, lstm.input_size)).swapaxes(0, 1)
        episode_starts = episode_starts.reshape((n_seq, -1)).swapaxes(0, 1)

        # If we don't have to reset the state in the middle of a sequence
        # we can avoid the for loop, which speeds up things
        if th.all(episode_starts == 0.0):
            lstm_output, lstm_states = lstm(features_sequence, lstm_states)
            lstm_output = th.flatten(lstm_output.transpose(0, 1), start_dim=0, end_dim=1)
            return lstm_output, lstm_states

        lstm_output = []
        # Iterate over the sequence
        for features, episode_start in zip(features_sequence, episode_starts, strict=True):
            hidden, lstm_states = lstm(
                features.unsqueeze(dim=0),
                (
                    # Reset the states at the beginning of a new episode
                    (1.0 - episode_start).view(1, n_seq, 1) * lstm_states[0],
                    (1.0 - episode_start).view(1, n_seq, 1) * lstm_states[1],
                ),
            )
            lstm_output += [hidden]
        # Sequence to batch
        # (sequence length, n_seq, lstm_out_dim) -> (batch_size, lstm_out_dim)
        lstm_output = th.flatten(th.cat(lstm_output).transpose(0, 1), start_dim=0, end_dim=1)
        return lstm_output, lstm_states
```

--------------------------------

### TQC _setup_model method

Source: https://sb3-contrib.readthedocs.io/en/master/_modules/sb3_contrib/tqc/tqc.html

Initializes the model by calling super()._setup_model(), creating aliases, and setting up batch norm statistics. Handles automatic target entropy and entropy coefficient learning, including validation of initial ent_coef values.

```python
    def _setup_model(self) -> None:
        super()._setup_model()
        self._create_aliases()
        # Running mean and running var
        self.batch_norm_stats = get_parameters_by_name(self.critic, ["running_"])
        self.batch_norm_stats_target = get_parameters_by_name(self.critic_target, ["running_"])

        # Target entropy is used when learning the entropy coefficient
        if self.target_entropy == "auto":
            # automatically set target entropy if needed
            self.target_entropy = -np.prod(self.env.action_space.shape).astype(np.float32)  # type: ignore
        else:
            # Force conversion
            # this will also throw an error for unexpected string
            self.target_entropy = float(self.target_entropy)

        # The entropy coefficient or entropy can be learned automatically
        # see Automating Entropy Adjustment for Maximum Entropy RL section
        # of https://arxiv.org/abs/1812.05905
        if isinstance(self.ent_coef, str) and self.ent_coef.startswith("auto"):
            # Default initial value of ent_coef when learned
            init_value = 1.0
            if "_" in self.ent_coef:
                init_value = float(self.ent_coef.split("_")[1])
                assert init_value > 0.0, "The initial value of ent_coef must be greater than 0"

            # Note: we optimize the log of the entropy coeff which is slightly different from the paper
            # as discussed in https://github.com/rail-berkeley/softlearning/issues/37
            self.log_ent_coef = th.log(th.ones(1, device=self.device) * init_value).requires_grad_(True)
            self.ent_coef_optimizer = th.optim.Adam([self.log_ent_coef], lr=self.lr_schedule(1))
        else:
            # Force conversion to float
            # this will throw an error if a malformed string (different from 'auto')
            # is passed
            self.ent_coef_tensor = th.tensor(float(self.ent_coef), device=self.device)
```

--------------------------------

### MaskablePPO Training and Evaluation Example

Source: https://sb3-contrib.readthedocs.io/en/master/modules/ppo_mask.html

Example demonstrating how to train a MaskablePPO agent on an environment with invalid action masks, evaluate it, save/load the model, and use action masks during prediction.

```APIDOC
## MaskablePPO Training and Evaluation Example

### Description
This example shows how to train a MaskablePPO agent on an environment with invalid action masks, evaluate it, save/load the model, and use action masks during prediction.

### Code
```python
from sb3_contrib import MaskablePPO
from sb3_contrib.common.envs import InvalidActionEnvDiscrete
from sb3_contrib.common.maskable.evaluation import evaluate_policy
from sb3_contrib.common.maskable.utils import get_action_masks
# This is a drop-in replacement for EvalCallback
from sb3_contrib.common.maskable.callbacks import MaskableEvalCallback


env = InvalidActionEnvDiscrete(dim=80, n_invalid_actions=60)
model = MaskablePPO("MlpPolicy", env, gamma=0.4, seed=32, verbose=1)
model.learn(5_000)

evaluate_policy(model, env, n_eval_episodes=20, reward_threshold=90, warn=False)

model.save("ppo_mask")
del model # remove to demonstrate saving and loading

model = MaskablePPO.load("ppo_mask")

obs, _ = env.reset()
while True:
    # Retrieve current action mask
    action_masks = get_action_masks(env)
    action, _states = model.predict(obs, action_masks=action_masks)
    obs, reward, terminated, truncated, info = env.step(action)
```
```

--------------------------------

### QRDQN __init__ method

Source: https://sb3-contrib.readthedocs.io/en/master/_modules/sb3_contrib/qrdqn/qrdqn.html

Initializes the QRDQN algorithm with parameters for policy, environment, learning rate, buffer size, exploration, and more. Sets up the parent OffPolicyAlgorithm with no action noise, SDE support disabled, and discrete action space. Configures exploration parameters and optimizer defaults (Adam with eps=0.01/batch_size) if not provided.

```python
    def __init__(
        self,
        policy: str | type[QRDQNPolicy],
        env: GymEnv | str,
        learning_rate: float | Schedule = 5e-5,
        buffer_size: int = 1000000,  # 1e6
        learning_starts: int = 100,
        batch_size: int = 32,
        tau: float = 1.0,
        gamma: float = 0.99,
        train_freq: int | tuple[int, str] = 4,
        gradient_steps: int = 1,
        replay_buffer_class: type[ReplayBuffer] | None = None,
        replay_buffer_kwargs: dict[str, Any] | None = None,
        optimize_memory_usage: bool = False,
        n_steps: int = 1,
        target_update_interval: int = 10000,
        exploration_fraction: float = 0.005,
        exploration_initial_eps: float = 1.0,
        exploration_final_eps: float = 0.01,
        max_grad_norm: float | None = None,
        stats_window_size: int = 100,
        tensorboard_log: str | None = None,
        policy_kwargs: dict[str, Any] | None = None,
        verbose: int = 0,
        seed: int | None = None,
        device: th.device | str = "auto",
        _init_setup_model: bool = True,
    ):
        super().__init__(
            policy,
            env,
            learning_rate,
            buffer_size,
            learning_starts,
            batch_size,
            tau,
            gamma,
            train_freq,
            gradient_steps,
            action_noise=None,  # No action noise
            replay_buffer_class=replay_buffer_class,
            replay_buffer_kwargs=replay_buffer_kwargs,
            optimize_memory_usage=optimize_memory_usage,
            n_steps=n_steps,
            policy_kwargs=policy_kwargs,
            stats_window_size=stats_window_size,
            tensorboard_log=tensorboard_log,
            verbose=verbose,
            device=device,
            seed=seed,
            sde_support=False,
            supported_action_spaces=(spaces.Discrete,),
            support_multi_env=True,
        )

        self.exploration_initial_eps = exploration_initial_eps
        self.exploration_final_eps = exploration_final_eps
        self.exploration_fraction = exploration_fraction
        self.target_update_interval = target_update_interval
        # For updating the target network with multiple envs:
        self._n_calls = 0
        self.max_grad_norm = max_grad_norm
        # "epsilon" for the epsilon-greedy exploration
        self.exploration_rate = 0.0

        if "optimizer_class" not in self.policy_kwargs:
            self.policy_kwargs["optimizer_class"] = th.optim.Adam
            # Proposed in the QR-DQN paper where `batch_size = 32`
            self.policy_kwargs["optimizer_kwargs"] = dict(eps=0.01 / batch_size)

        if _init_setup_model:
            self._setup_model()
```

--------------------------------

### QRDQNPolicy _build method

Source: https://sb3-contrib.readthedocs.io/en/master/_modules/sb3_contrib/qrdqn/policies.html

Creates the quantile network and its target network, loads the target's state from the main network, and sets up the optimizer with the initial learning rate.

```python
    def _build(self, lr_schedule: Schedule) -> None:
        """
        Create the network and the optimizer.

        :param lr_schedule: Learning rate schedule
            lr_schedule(1) is the initial learning rate
        """
        self.quantile_net = self.make_quantile_net()
        self.quantile_net_target = self.make_quantile_net()
        self.quantile_net_target.load_state_dict(self.quantile_net.state_dict())
        self.quantile_net_target.set_training_mode(False)

        # Setup optimizer with initial learning rate
        self.optimizer = self.optimizer_class(  # type: ignore[call-arg]
            self.quantile_net.parameters(),
            lr=lr_schedule(1),
            **self.optimizer_kwargs,
        )
```

--------------------------------

### ActorCriticPolicy.get_distribution

Source: https://sb3-contrib.readthedocs.io/en/master/_modules/stable_baselines3/common/policies.html

Gets the current policy distribution given the observations.

```APIDOC
## ActorCriticPolicy.get_distribution

### Description
Gets the current policy distribution given the observations.

### Method
`get_distribution(obs)`

### Parameters
- **obs** (PyTorchObs) - Required - Observation.

### Returns
- The action distribution.
```

--------------------------------

### RecurrentActorCriticPolicy.get_distribution

Source: https://sb3-contrib.readthedocs.io/en/master/modules/ppo_recurrent.html

Get the current policy distribution given the observations.

```APIDOC
## RecurrentActorCriticPolicy.get_distribution

### Description
Get the current policy distribution given the observations.

### Method
get_distribution

### Parameters
- **obs** (_Tensor_) - Observation.
- **lstm_states** (_tuple_[Tensor, Tensor]) - The last hidden and memory states for the LSTM.
- **episode_starts** (_Tensor_) - Whether the observations correspond to new episodes or not (we reset the lstm states in that case).

### Returns
- **tuple** (_Distribution, tuple[Tensor, ...]) - the action distribution and new hidden states.
```

--------------------------------

### MaskablePPO._setup_model method

Source: https://sb3-contrib.readthedocs.io/en/master/_modules/sb3_contrib/ppo_mask/ppo_mask.html

Sets up the model by initializing the learning rate schedule, random seed, policy, and rollout buffer. Validates that the policy is a MaskableActorCriticPolicy and selects the appropriate maskable rollout buffer for Dict or non-Dict observation spaces. Also initializes clipping schedules, requiring clip_range_vf to be positive if provided.

```python
    def _setup_model(self) -> None:
        self._setup_lr_schedule()
        self.set_random_seed(self.seed)

        self.policy = self.policy_class(  # type: ignore[assignment]
            self.observation_space,
            self.action_space,
            self.lr_schedule,
            **self.policy_kwargs,
        )
        self.policy = self.policy.to(self.device)

        if not isinstance(self.policy, MaskableActorCriticPolicy):
            raise ValueError("Policy must subclass MaskableActorCriticPolicy")

        if self.rollout_buffer_class is None:
            if isinstance(self.observation_space, spaces.Dict):
                self.rollout_buffer_class = MaskableDictRolloutBuffer
            else:
                self.rollout_buffer_class = MaskableRolloutBuffer

        self.rollout_buffer = self.rollout_buffer_class(  # type: ignore[assignment]
            self.n_steps,
            self.observation_space,
            self.action_space,
            self.device,
            gamma=self.gamma,
            gae_lambda=self.gae_lambda,
            n_envs=self.n_envs,
            **self.rollout_buffer_kwargs,
        )

        # Initialize schedules for policy/value clipping
        self.clip_range = FloatSchedule(self.clip_range)
        if self.clip_range_vf is not None:
            if isinstance(self.clip_range_vf, (float, int)):
                assert self.clip_range_vf > 0, "`clip_range_vf` must be positive, " "pass `None` to deactivate vf clipping"

            self.clip_range_vf = FloatSchedule(self.clip_range_vf)
```

--------------------------------

### ActorCriticPolicy.predict_values

Source: https://sb3-contrib.readthedocs.io/en/master/_modules/stable_baselines3/common/policies.html

Gets the estimated values according to the current policy given the observations.

```APIDOC
## ActorCriticPolicy.predict_values

### Description
Gets the estimated values according to the current policy given the observations.

### Method
`predict_values(obs)`

### Parameters
- **obs** (PyTorchObs) - Required - Observation.

### Returns
- The estimated values.
```

--------------------------------

### BaseAlgorithm __init__ method

Source: https://sb3-contrib.readthedocs.io/en/master/_modules/stable_baselines3/common/base_class.html

Initializes the algorithm with policy, environment, learning rate, and other settings. Validates action spaces, multi-env support, and gSDE usage. Raises ValueError for unsupported configurations.

```python
    def __init__(
        self,
        policy: str | type[BasePolicy],
        env: GymEnv | str | None,
        learning_rate: float | Schedule,
        policy_kwargs: dict[str, Any] | None = None,
        stats_window_size: int = 100,
        tensorboard_log: str | None = None,
        verbose: int = 0,
        device: th.device | str = "auto",
        support_multi_env: bool = False,
        monitor_wrapper: bool = True,
        seed: int | None = None,
        use_sde: bool = False,
        sde_sample_freq: int = -1,
        supported_action_spaces: tuple[type[spaces.Space], ...] | None = None,
    ) -> None:
        if isinstance(policy, str):
            self.policy_class = self._get_policy_from_name(policy)
        else:
            self.policy_class = policy

        self.device = get_device(device)
        if verbose >= 1:
            print(f"Using {self.device} device")

        self.verbose = verbose
        self.policy_kwargs = {} if policy_kwargs is None else policy_kwargs

        self.num_timesteps = 0
        # Used for updating schedules
        self._total_timesteps = 0
        # Used for computing fps, it is updated at each call of learn()
        self._num_timesteps_at_start = 0
        self.seed = seed
        self.action_noise: ActionNoise | None = None
        self.start_time = 0.0
        self.learning_rate = learning_rate
        self.tensorboard_log = tensorboard_log
        self._last_obs = None  # type: np.ndarray | dict[str, np.ndarray] | None
        self._last_episode_starts = None  # type: np.ndarray | None
        # When using VecNormalize:
        self._last_original_obs = None  # type: np.ndarray | dict[str, np.ndarray] | None
        self._episode_num = 0
        # Used for gSDE only
        self.use_sde = use_sde
        self.sde_sample_freq = sde_sample_freq
        # Track the training progress remaining (from 1 to 0)
        # this is used to update the learning rate
        self._current_progress_remaining = 1.0
        # Buffers for logging
        self._stats_window_size = stats_window_size
        self.ep_info_buffer = None  # type: deque | None
        self.ep_success_buffer = None  # type: deque | None
        # For logging (and TD3 delayed updates)
        self._n_updates = 0  # type: int
        # Whether the user passed a custom logger or not
        self._custom_logger = False
        self.env: VecEnv | None = None
        self._vec_normalize_env: VecNormalize | None = None

        # Create and wrap the env if needed
        if env is not None:
            env = maybe_make_env(env, self.verbose)
            env = self._wrap_env(env, self.verbose, monitor_wrapper)

            self.observation_space = env.observation_space
            self.action_space = env.action_space
            self.n_envs = env.num_envs
            self.env = env

            # get VecNormalize object if needed
            self._vec_normalize_env = unwrap_vec_normalize(env)

            if supported_action_spaces is not None:
                assert isinstance(self.action_space, supported_action_spaces), (
                    f"The algorithm only supports {supported_action_spaces} as action spaces "
                    f"but {self.action_space} was provided"
                )

            if not support_multi_env and self.n_envs > 1:
                raise ValueError(
                    "Error: the model does not support multiple envs; it requires " "a single vectorized environment."
                )

            # Catch common mistake: using MlpPolicy/CnnPolicy instead of MultiInputPolicy
            if policy in ["MlpPolicy", "CnnPolicy"] and isinstance(self.observation_space, spaces.Dict):
                raise ValueError(f"You must use `MultiInputPolicy` when working with dict observation space, not {policy}")

            if self.use_sde and not isinstance(self.action_space, spaces.Box):
                raise ValueError("generalized State-Dependent Exploration (gSDE) can only be used with continuous actions.")

            if isinstance(self.action_space, spaces.Box):
                assert np.all(
                    np.isfinite(np.array([self.action_space.low, self.action_space.high]))
                ), "Continuous action space must have a finite lower and upper bound"
```

--------------------------------

### CrossQ Actor: action distribution setup and constructor parameters

Source: https://sb3-contrib.readthedocs.io/en/master/_modules/sb3_contrib/crossq/policies.html

Sets up the action distribution for the CrossQ actor, either using StateDependentNoiseDistribution (gSDE) or SquashedDiagGaussianDistribution. Also collects constructor parameters for serialization.

```python
        if self.use_sde:
            self.action_dist = StateDependentNoiseDistribution(
                action_dim, full_std=full_std, use_expln=use_expln, learn_features=True, squash_output=True
            )
            self.mu, self.log_std = self.action_dist.proba_distribution_net(
                latent_dim=last_layer_dim, latent_sde_dim=last_layer_dim, log_std_init=log_std_init
            )
            # Avoid numerical issues by limiting the mean of the Gaussian
            # to be in [-clip_mean, clip_mean]
            if clip_mean > 0.0:
                self.mu = nn.Sequential(self.mu, nn.Hardtanh(min_val=-clip_mean, max_val=clip_mean))
        else:
            self.action_dist = SquashedDiagGaussianDistribution(action_dim)  # type: ignore[assignment]
            self.mu = nn.Linear(last_layer_dim, action_dim)
            self.log_std = nn.Linear(last_layer_dim, action_dim)  # type: ignore[assignment]
```

```python
    def _get_constructor_parameters(self) -> dict[str, Any]:
        data = super()._get_constructor_parameters()

        data.update(
            dict(
                net_arch=self.net_arch,
                features_dim=self.features_dim,
                activation_fn=self.activation_fn,
                use_sde=self.use_sde,
                log_std_init=self.log_std_init,
                full_std=self.full_std,
                use_expln=self.use_expln,
                features_extractor=self.features_extractor,
                clip_mean=self.clip_mean,
            )
        )
        return data
```

--------------------------------

### predict_values

Source: https://sb3-contrib.readthedocs.io/en/master/modules/ppo_mask.html

Get the estimated values according to the current policy given the observations.

```APIDOC
## predict_values

### Description
Get the estimated values according to the current policy given the observations.

### Method
Function

### Signature
`predict_values(obs)`

### Parameters
- **obs** (Tensor | dict[str, Tensor]) - Required - Observation

### Returns
- **Tensor** - the estimated values.
```

--------------------------------

### RecurrentPPO _setup_model method

Source: https://sb3-contrib.readthedocs.io/en/master/_modules/sb3_contrib/ppo_recurrent/ppo_recurrent.html

Sets up the model for RecurrentPPO, including the policy, LSTM state buffers, and rollout buffer. Validates that the policy is a RecurrentActorCriticPolicy.

```python
    def _setup_model(self) -> None:
        self._setup_lr_schedule()
        self.set_random_seed(self.seed)

        buffer_cls = RecurrentDictRolloutBuffer if isinstance(self.observation_space, spaces.Dict) else RecurrentRolloutBuffer

        self.policy = self.policy_class(
            self.observation_space,
            self.action_space,
            self.lr_schedule,
            use_sde=self.use_sde,
            **self.policy_kwargs,
        )
        self.policy = self.policy.to(self.device)

        # We assume that LSTM for the actor and the critic
        # have the same architecture
        lstm = self.policy.lstm_actor

        if not isinstance(self.policy, RecurrentActorCriticPolicy):
            raise ValueError("Policy must subclass RecurrentActorCriticPolicy")

        single_hidden_state_shape = (lstm.num_layers, self.n_envs, lstm.hidden_size)
        # hidden and cell states for actor and critic
        self._last_lstm_states = RNNStates(
            (
                th.zeros(single_hidden_state_shape, device=self.device),
                th.zeros(single_hidden_state_shape, device=self.device),
            ),
            (
                th.zeros(single_hidden_state_shape, device=self.device),
                th.zeros(single_hidden_state_shape, device=self.device),
            ),
        )

        hidden_state_buffer_shape = (self.n_steps, lstm.num_layers, self.n_envs, lstm.hidden_size)

        self.rollout_buffer = buffer_cls(
            self.n_steps,
            self.observation_space,
            self.action_space,
            hidden_state_buffer_shape,
            self.device,
            gamma=self.gamma,
            gae_lambda=self.gae_lambda,
            n_envs=self.n_envs,
        )

        # Initialize schedules for policy/value clipping
        self.clip_range = FloatSchedule(self.clip_range)
        if self.clip_range_vf is not None:
            if isinstance(self.clip_range_vf, (float, int)):
                assert self.clip_range_vf > 0, "`clip_range_vf` must be positive, pass `None` to deactivate vf clipping"

            self.clip_range_vf = FloatSchedule(self.clip_range_vf)
```

--------------------------------

### RecurrentActorCriticPolicy.predict_values

Source: https://sb3-contrib.readthedocs.io/en/master/modules/ppo_recurrent.html

Get the estimated values according to the current policy given the observations.

```APIDOC
## RecurrentActorCriticPolicy.predict_values

### Description
Get the estimated values according to the current policy given the observations.

### Method
predict_values

### Parameters
- **obs** (_Tensor_) - Observation.
- **lstm_states** (_tuple_[Tensor, Tensor]) - The last hidden and memory states for the LSTM.
- **episode_starts** (_Tensor_) - Whether the observations correspond to new episodes or not (we reset the lstm states in that case).

### Returns
- **_Tensor** - the estimated values.
```

--------------------------------

### MaskableActorCriticPolicy.predict_values

Source: https://sb3-contrib.readthedocs.io/en/master/_modules/sb3_contrib/common/maskable/policies.html

Gets the estimated values according to the current policy given the observations.

```APIDOC
## predict_values

### Description
Gets the estimated values according to the current policy given the observations.

### Method
predict_values

### Parameters
- **obs** (PyTorchObs) - Observation

### Returns
- **values** (th.Tensor) - The estimated values
```

--------------------------------

### TQC __init__ method

Source: https://sb3-contrib.readthedocs.io/en/master/_modules/sb3_contrib/tqc/tqc.html

Constructor for the TQC class. Sets up the algorithm with hyperparameters such as learning rate, buffer size, and top_quantiles_to_drop_per_net. Calls super().__init__ with supported action spaces of Box and multi-environment support.

```python
    def __init__(
        self,
        policy: str | type[TQCPolicy],
        env: GymEnv | str,
        learning_rate: float | Callable = 3e-4,
        buffer_size: int = 1000000,  # 1e6
        learning_starts: int = 100,
        batch_size: int = 256,
        tau: float = 0.005,
        gamma: float = 0.99,
        train_freq: int | tuple[int, str] = 1,
        gradient_steps: int = 1,
        action_noise: ActionNoise | None = None,
        replay_buffer_class: type[ReplayBuffer] | None = None,
        replay_buffer_kwargs: dict[str, Any] | None = None,
        optimize_memory_usage: bool = False,
        n_steps: int = 1,
        ent_coef: str | float = "auto",
        target_update_interval: int = 1,
        target_entropy: str | float = "auto",
        top_quantiles_to_drop_per_net: int = 2,
        use_sde: bool = False,
        sde_sample_freq: int = -1,
        use_sde_at_warmup: bool = False,
        stats_window_size: int = 100,
        tensorboard_log: str | None = None,
        policy_kwargs: dict[str, Any] | None = None,
        verbose: int = 0,
        seed: int | None = None,
        device: th.device | str = "auto",
        _init_setup_model: bool = True,
    ):
        super().__init__(
            policy,
            env,
            learning_rate,
            buffer_size,
            learning_starts,
            batch_size,
            tau,
            gamma,
            train_freq,
            gradient_steps,
            action_noise=action_noise,
            replay_buffer_class=replay_buffer_class,
            replay_buffer_kwargs=replay_buffer_kwargs,
            optimize_memory_usage=optimize_memory_usage,
            n_steps=n_steps,
            policy_kwargs=policy_kwargs,
            stats_window_size=stats_window_size,
            tensorboard_log=tensorboard_log,
            verbose=verbose,
            device=device,
            seed=seed,
            use_sde=use_sde,
            sde_sample_freq=sde_sample_freq,
            use_sde_at_warmup=use_sde_at_warmup,
            supported_action_spaces=(spaces.Box,),
            support_multi_env=True,
        )

        self.target_entropy = target_entropy
        self.log_ent_coef = None  # type: th.Tensor | None
        # Entropy coefficient / Entropy temperature
        # Inverse of the reward scale
        self.ent_coef = ent_coef
        self.target_update_interval = target_update_interval
        self.ent_coef_optimizer: th.optim.Adam | None = None
        self.top_quantiles_to_drop_per_net = top_quantiles_to_drop_per_net

        if _init_setup_model:
            self._setup_model()
```

--------------------------------

### ActorCriticPolicy._predict

Source: https://sb3-contrib.readthedocs.io/en/master/_modules/stable_baselines3/common/policies.html

Gets the action according to the policy for a given observation, either stochastic or deterministic.

```APIDOC
## ActorCriticPolicy._predict

### Description
Gets the action according to the policy for a given observation, either stochastic or deterministic.

### Method
`_predict(observation, deterministic=False)`

### Parameters
- **observation** (PyTorchObs) - Required - The observation.
- **deterministic** (bool) - Optional - Whether to use stochastic or deterministic actions.

### Returns
- Taken action according to the policy.
```

--------------------------------

### QRDQN usage example with CartPole

Source: https://sb3-contrib.readthedocs.io/en/master/modules/qrdqn.html

Shows how to create, train, save, load, and use a QRDQN model with the CartPole-v1 environment. Includes policy_kwargs with n_quantiles=50 and a prediction loop with deterministic actions.

```python
import gymnasium as gym

from sb3_contrib import QRDQN

env = gym.make("CartPole-v1", render_mode="human")

policy_kwargs = dict(n_quantiles=50)
model = QRDQN("MlpPolicy", env, policy_kwargs=policy_kwargs, verbose=1)
model.learn(total_timesteps=10_000, log_interval=4)
model.save("qrdqn_cartpole")

del model # remove to demonstrate saving and loading

model = QRDQN.load("qrdqn_cartpole")

obs, _ = env.reset()
while True:
    action, _states = model.predict(obs, deterministic=True)
    obs, reward, terminated, truncated, info = env.step(action)
    env.render()
    if terminated or truncated:
        obs, _ = env.reset()
```

--------------------------------

### get_distribution

Source: https://sb3-contrib.readthedocs.io/en/master/modules/ppo_mask.html

Get the current policy distribution given the observations and optional action masks.

```APIDOC
## get_distribution

### Description
Get the current policy distribution given the observations.

### Method
Function

### Signature
`get_distribution(obs, action_masks=None)`

### Parameters
- **obs** (Tensor | dict[str, Tensor]) - Required - Observation
- **action_masks** (ndarray | None) - Optional - Actions' mask

### Returns
- **MaskableDistribution** - the action distribution
```

--------------------------------

### Get current environment

Source: https://sb3-contrib.readthedocs.io/en/master/_modules/stable_baselines3/common/base_class.html

Returns the current environment, which can be None if not defined. Part of the BaseAlgorithm API.

```python
    def get_env(self) -> VecEnv | None:
        """
        Returns the current environment (can be None if not defined).

        :return: The current environment
        """
        return self.env
```

--------------------------------

### MaskableActorCriticPolicy.get_distribution

Source: https://sb3-contrib.readthedocs.io/en/master/_modules/sb3_contrib/common/maskable/policies.html

Gets the current policy distribution given the observations. Applies action masking if provided.

```APIDOC
## get_distribution

### Description
Gets the current policy distribution given the observations. Applies action masking if provided.

### Method
get_distribution

### Parameters
- **obs** (PyTorchObs) - Observation
- **action_masks** (np.ndarray | None) - Optional action masks

### Returns
- **distribution** (MaskableDistribution) - The action distribution
```

--------------------------------

### TQCPolicy critic setup and optimizer

Source: https://sb3-contrib.readthedocs.io/en/master/_modules/sb3_contrib/tqc/policies.html

Configures the critic and critic target networks, handling shared or separate feature extractors, and creates the critic optimizer with the given learning rate schedule.

```python
        if self.share_features_extractor:
            self.critic = self.make_critic(features_extractor=self.actor.features_extractor)
            # Do not optimize the shared features extractor with the critic loss
            # otherwise, there are gradient computation issues
            critic_parameters = [param for name, param in self.critic.named_parameters() if "features_extractor" not in name]
        else:
            # Create a separate features extractor for the critic
            # this requires more memory and computation
            self.critic = self.make_critic(features_extractor=None)
            critic_parameters = list(self.critic.parameters())

        # Critic target should not share the feature extactor with critic
        self.critic_target = self.make_critic(features_extractor=None)
        self.critic_target.load_state_dict(self.critic.state_dict())

        # Target networks should always be in eval mode
        self.critic_target.set_training_mode(False)

        self.critic.optimizer = self.optimizer_class(  # type: ignore[call-arg]
            critic_parameters,
            lr=lr_schedule(1),
            **self.optimizer_kwargs,
        )
```
Tokens:
Raw