### _setup_learn

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

Sets up the necessary variables and environment configurations before starting the learning process.

```APIDOC
## _setup_learn

### Description
Initializes different variables needed for training.

### Parameters
#### Parameters
- **total_timesteps** (int) - The total number of samples (env steps) to train on
- **callback** (MaybeCallback) - Optional - Callback(s) called at every step with state of the algorithm.
- **reset_num_timesteps** (bool) - Optional - Whether to reset or not the ``num_timesteps`` attribute. Defaults to True.
- **tb_log_name** (str) - Optional - the name of the run for tensorboard log. Defaults to "run".
- **progress_bar** (bool) - Optional - Display a progress bar using tqdm and rich.

### Returns
- **tuple[int, BaseCallback]** - Total timesteps and callback(s)
```

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

### Install Development Version

Source: https://sb3-contrib.readthedocs.io/en/master/_sources/guide/install.md.txt

Clone the repository and install the library in editable mode for development. This includes support for running tests and building documentation.

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

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

### Model Setup for QR-DQN

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 for network components, and initializing the exploration schedule. It also includes a warning for specific environment and target update configurations.

```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."
                )

```

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

### Install Dependencies for MicroRTS and Project

Source: https://sb3-contrib.readthedocs.io/en/master/_sources/modules/ppo_mask.md.txt

Install MicroRTS by downloading and extracting it, then install the project's Python dependencies. It is recommended to use a virtual environment.

```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
```

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

### TQC Usage Example

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

This example demonstrates how to instantiate, train, and use a TQC model with a custom policy and environment.

```APIDOC
## TQC Usage Example

### Description
This example shows how to create a TQC agent, train it on an environment, save it, and then load it for inference.

### Method
```python
import gymnasium as gym
import numpy as np

from sb3_contrib import TQC

# Instantiate the environment
env = gym.make("Pendulum-v1", render_mode="human")

# Define policy keyword arguments for TQC
policy_kwargs = dict(n_critics=2, n_quantiles=25)

# Create a TQC model with MlpPolicy
model = TQC("MlpPolicy", env, top_quantiles_to_drop_per_net=2, verbose=1, policy_kwargs=policy_kwargs)

# Train the model
model.learn(total_timesteps=10_000, log_interval=4)

# Save the trained model
model.save("tqc_pendulum")

# Remove the model to demonstrate loading
del model

# Load the saved model
model = TQC.load("tqc_pendulum")

# Perform inference
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()
```

### Parameters
- `env`: The gymnasium environment to train on.
- `policy_kwargs`: Dictionary for policy-specific parameters like `n_critics` and `n_quantiles`.
- `top_quantiles_to_drop_per_net`: Number of quantiles to drop per network.
- `verbose`: Verbosity level (0 for silent, 1 for progress bar, 2 for detailed logging).

### Request Example
```python
# See the Method section for a complete example.
```

### Response
- The `learn` method returns `None`.
- The `predict` method returns the action and the next state.
```

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

### QR-DQN Usage Example

Source: https://sb3-contrib.readthedocs.io/en/master/_sources/modules/qrdqn.md.txt

This example demonstrates how to instantiate, train, and use a QR-DQN model with a Gymnasium environment.

```APIDOC
## QR-DQN Usage Example

### Description
This code snippet shows how to initialize a QR-DQN agent, train it on a given environment, save the trained model, and then load it for inference.

### Method
```python
import gymnasium as gym

from sb3_contrib import QRDQN

# Initialize the environment
env = gym.make("CartPole-v1", render_mode="human")

# Define policy keyword arguments, e.g., number of quantiles
policy_kwargs = dict(n_quantiles=50)

# Instantiate the QR-DQN model
model = QRDQN("MlpPolicy", env, policy_kwargs=policy_kwargs, verbose=1)

# Train the model
model.learn(total_timesteps=10_000, log_interval=4)

# Save the model
model.save("qrdqn_cartpole")

# Load the model
del model  # remove to demonstrate saving and loading
model = QRDQN.load("qrdqn_cartpole")

# Perform inference
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()
```

### Parameters
- `env`: The Gymnasium environment to train on.
- `policy_kwargs`: Dictionary of additional arguments for the policy, such as `n_quantiles`.
- `verbose`: Verbosity level (0 for silent, 1 for progress bar, 2 for detailed logging).

### Training
- `total_timesteps`: The total number of samples (timesteps) to train for.
- `log_interval`: Log progress every `log_interval` updates.

### Inference
- `predict`: Method to get the action from the current observation.
  - `obs`: The current observation from the environment.
  - `deterministic`: If True, use the most likely action; otherwise, sample from the action distribution.
```

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

### CrossQ Usage Example

Source: https://sb3-contrib.readthedocs.io/en/master/_sources/modules/crossq.md.txt

A basic example demonstrating how to instantiate and train a CrossQ model with an MLP policy.

```python
from sb3_contrib import CrossQ

# Instantiate the model with MlpPolicy for Walker2d-v4 environment
model = CrossQ("MlpPolicy", "Walker2d-v4")

# Train the model for a specified number of timesteps
model.learn(total_timesteps=1_000_000)

# Save the trained model
model.save("crossq_walker")
```

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

### CrossQ Model Setup

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

Sets up the CrossQ model, including calculating the target entropy and initializing the entropy coefficient optimizer if it's set to 'auto'.

```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)

```

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

### Install Stable Release

Source: https://sb3-contrib.readthedocs.io/en/master/_sources/guide/install.md.txt

Install the latest stable version of Stable-Baselines3 contrib using pip.

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

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

### Train and Load TQC Model

Source: https://sb3-contrib.readthedocs.io/en/master/_sources/modules/tqc.md.txt

Example of training a TQC model with specified policy kwargs, saving it, and then loading it for inference. Ensure 'Pendulum-v1' environment is installed.

```python
import gymnasium as gym
import numpy as np

from sb3_contrib import TQC

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

policy_kwargs = dict(n_critics=2, n_quantiles=25)
model = TQC("MlpPolicy", env, top_quantiles_to_drop_per_net=2, verbose=1, policy_kwargs=policy_kwargs)
model.learn(total_timesteps=10_000, log_interval=4)
model.save("tqc_pendulum")

del model # remove to demonstrate saving and loading

model = TQC.load("tqc_pendulum")

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()

```

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

### TRPO Usage Example

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

This example demonstrates how to instantiate, train, and use a TRPO model with a Gymnasium environment. It covers saving and loading the model as well.

```APIDOC
## TRPO Example

```python
import gymnasium as gym
import numpy as np

from sb3_contrib import TRPO

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

model = TRPO("MlpPolicy", env, verbose=1)
model.learn(total_timesteps=10_000, log_interval=4)
model.save("trpo_pendulum")

del model # remove to demonstrate saving and loading

model = TRPO.load("trpo_pendulum")

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()

```
```

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

### Install Bleeding-edge Version

Source: https://sb3-contrib.readthedocs.io/en/master/_sources/guide/install.md.txt

Install the most recent version directly from the GitHub repository using pip.

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

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

### RecurrentPPO Example

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

This example demonstrates how to instantiate, train, and use the RecurrentPPO model. It includes saving and loading the model, and importantly, how to handle LSTM states during prediction.

```APIDOC
## RecurrentPPO Usage Example

### Description
This example demonstrates the basic usage of the `RecurrentPPO` class, including model instantiation, training, evaluation, saving, loading, and prediction with LSTM state management.

### Method
```python
import numpy as np

from sb3_contrib import RecurrentPPO
from stable_baselines3.common.evaluation import evaluate_policy

# Instantiate the model with a recurrent policy
model = RecurrentPPO("MlpLstmPolicy", "CartPole-v1", verbose=1)

# Train the model
model.learn(5000)

# Evaluate the trained model
vec_env = model.get_env()
mean_reward, std_reward = evaluate_policy(model, vec_env, n_eval_episodes=20, warn=False)
print(f"Mean reward: {mean_reward} +/- {std_reward}")

# Save the model
model.save("ppo_recurrent")
del model # remove to demonstrate saving and loading

# Load the model
model = RecurrentPPO.load("ppo_recurrent")

# Predict with LSTM states
obs = vec_env.reset()
# Initialize LSTM states and episode start signals
lstm_states = None
num_envs = 1
episode_starts = np.ones((num_envs,), dtype=bool)

while True:
    # Predict action, updating LSTM states
    action, lstm_states = model.predict(obs, state=lstm_states, episode_start=episode_starts, deterministic=True)
    obs, rewards, dones, info = vec_env.step(action)
    # Update episode start signals for the next step
    episode_starts = dones
    vec_env.render("human")

```

### Parameters
- `policy` (str or `RecurrentActorCriticPolicy`): The policy to use (e.g., "MlpLstmPolicy", "CnnLstmPolicy").
- `env` (str or `gym.Env`): The environment to train on.
- `verbose` (int): Verbosity level (0 for silent, 1 for progress bar, 2 for detailed logging).

### Request Body
N/A for instantiation.

### Response
N/A for instantiation.

### Notes
- It is crucial to pass the `lstm_states` and `episode_start` arguments to the `predict()` method to ensure correct LSTM state updates.
- Available policies include `MlpLstmPolicy`, `CnnLstmPolicy`, and `MultiInputLstmPolicy`.
```

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

### QR-DQN Example with CartPole

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

Demonstrates how to instantiate, train, and use a QR-DQN model with the CartPole environment. Includes saving and loading the model. Requires gymnasium and sb3_contrib.

```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()

```

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

### Recurrent PPO Example with LSTM States

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

Demonstrates training, saving, loading, and prediction with RecurrentPPO. It is crucial to pass `lstm_states` and `episode_start` to `predict()` for correct LSTM state updates. This example uses `MlpLstmPolicy` for a CartPole environment.

```python
import numpy as np

from sb3_contrib import RecurrentPPO
from stable_baselines3.common.evaluation import evaluate_policy

model = RecurrentPPO("MlpLstmPolicy", "CartPole-v1", verbose=1)
model.learn(5000)

vec_env = model.get_env()
mean_reward, std_reward = evaluate_policy(model, vec_env, n_eval_episodes=20, warn=False)
print(mean_reward)

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

model = RecurrentPPO.load("ppo_recurrent")

obs = vec_env.reset()
# cell and hidden state of the LSTM
lstm_states = None
num_envs = 1
# Episode start signals are used to reset the lstm states
episode_starts = np.ones((num_envs,), dtype=bool)
while True:
    action, lstm_states = model.predict(obs, state=lstm_states, episode_start=episode_starts, deterministic=True)
    obs, rewards, dones, info = vec_env.step(action)
    episode_starts = dones
    vec_env.render("human")

```

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

### Training PPO with Masking

Source: https://sb3-contrib.readthedocs.io/en/master/_sources/modules/ppo_mask.md.txt

Instructions on how to train PPO agents with and without masking for different environment sizes, including example commands.

```APIDOC
## Training PPO with Masking

This section provides commands to replicate experiment results by training PPO agents with and without masking.

### Usage

```bash
# python sb3/train_ppo.py [output dir] [MicroRTS map size] [--mask] [--seed int]
```

### Examples

#### 4x4 Environment

*   **Unmasked:**
    ```bash
    python sb3/train_ppo.py zoo 4 --seed 42
    python sb3/train_ppo.py zoo 4 --seed 43
    python sb3/train_ppo.py zoo 4 --seed 44
    ```
*   **Masked:**
    ```bash
    python sb3/train_ppo.py zoo 4 --mask --seed 42
    python sb3/train_ppo.py zoo 4 --mask --seed 43
    python sb3/train_ppo.py zoo 4 --mask --seed 44
    ```

#### 10x10 Environment

*   **Unmasked:**
    ```bash
    python sb3/train_ppo.py zoo 10 --seed 42
    python sb3/train_ppo.py zoo 10 --seed 43
    python sb3/train_ppo.py zoo 10 --seed 44
    ```
*   **Masked:**
    ```bash
    python sb3/train_ppo.py zoo 10 --mask --seed 42
    python sb3/train_ppo.py zoo 10 --mask --seed 43
    python sb3/train_ppo.py zoo 10 --mask --seed 44
    ```

### Viewing Logs

Use TensorBoard to view the log output:

```bash
# For 4x4 environment
 tensorboard --logdir zoo/4x4/runs

# For 10x10 environment
 tensorboard --logdir zoo/10x10/runs
```
```

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

### Train PPO Models (10x10 Map)

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

Train PPO models for the 10x10 MicroRTS environment. Similar to the 4x4 examples, this demonstrates training with and without masking, using specified seeds.

```bash
# 10x10 unmasked
python sb3/train_ppo.py zoo 10 --seed 42
python sb3/train_ppo.py zoo 10 --seed 43
python sb3/train_ppo.py zoo 10 --seed 44

# 10x10 masked
python sb3/train_ppo.py zoo 10 --mask --seed 42
python sb3/train_ppo.py zoo 10 --mask --seed 43
python sb3/train_ppo.py zoo 10 --mask --seed 44
```

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

### RecurrentPPO Usage Example

Source: https://sb3-contrib.readthedocs.io/en/master/_sources/modules/ppo_recurrent.md.txt

Demonstrates how to instantiate, train, predict with, save, and load a RecurrentPPO model. It highlights the importance of passing `lstm_states` and `episode_start` to the `predict` method for proper state management.

```APIDOC
## RecurrentPPO Usage Example

This example shows the typical workflow for using the `RecurrentPPO` class.

### Method
```python
RecurrentPPO("MlpLstmPolicy", "CartPole-v1", verbose=1)
```

### Training
```python
model.learn(5000)
```

### Evaluation
```python
vec_env = model.get_env()
mean_reward, std_reward = evaluate_policy(model, vec_env, n_eval_episodes=20, warn=False)
print(mean_reward)
```

### Saving and Loading
```python
model.save("ppo_recurrent")
del model # remove to demonstrate saving and loading
model = RecurrentPPO.load("ppo_recurrent")
```

### Prediction with LSTM States
```python
obs = vec_env.reset()
lstm_states = None
num_envs = 1
episode_starts = np.ones((num_envs,), dtype=bool)
while True:
    action, lstm_states = model.predict(obs, state=lstm_states, episode_start=episode_starts, deterministic=True)
    obs, rewards, dones, info = vec_env.step(action)
    episode_starts = dones
    vec_env.render("human")
```
```

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

### Train PPO Models (4x4 Map)

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

Train PPO models for the 4x4 MicroRTS environment. Examples show training both with and without masking, using different seeds for reproducibility.

```bash
# 4x4 unmasked
python sb3/train_ppo.py zoo 4 --seed 42
python sb3/train_ppo.py zoo 4 --seed 43
python sb3/train_ppo.py zoo 4 --seed 44

# 4x4 masked
python sb3/train_ppo.py zoo 4 --mask --seed 42
python sb3/train_ppo.py zoo 4 --mask --seed 43
python sb3/train_ppo.py zoo 4 --mask --seed 44
```

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

### predict()

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

Gets the policy action for a given observation, optionally handling hidden states and episode starts. Supports deterministic or stochastic actions.

```APIDOC
## predict(_observation_ , _state =None_, _episode_start =None_, _deterministic =False_)

### Description
Get the policy action from an observation (and optional hidden state). Includes sugar-coating to handle different observations (e.g. normalizing images).

### Parameters
* **observation** (_type_) – The observation for which to get the action.
* **state** (_type_ | None) – The current state of the recurrent neural network (if any).
* **episode_start** (_type_ | None) – Flag to indicate if the current observation is the start of a new episode.
* **deterministic** (_bool_) – Whether to use deterministic or stochastic actions.
```

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

### _setup_model

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

Sets up the model by creating aliases for policy networks and initializing the exploration schedule.

```APIDOC
## _setup_model

### Description
Internal method to set up the model. It calls the parent class's setup, creates aliases for the quantile networks, copies batch normalization statistics, and initializes the exploration schedule.

### Method
`_setup_model()`
```

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

### Policy Prediction Method

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

Provides the main interface for getting actions from the policy, handling observation preprocessing, hidden states, and episode starts. It also includes a check for a common user error mixing Gym and SB3 VecEnv APIs.

```python
    def predict(
        self,
        observation: np.ndarray | dict[str, np.ndarray],
        state: tuple[np.ndarray, ...] | None = None,
        episode_start: np.ndarray | None = None,
        deterministic: bool = False,
    ) -> tuple[np.ndarray, tuple[np.ndarray, ...] | None]:
        """
        Get the policy action from an observation (and optional hidden state).
        Includes sugar-coating to handle different observations (e.g. normalizing images).

        :param observation: the input observation
        :param state: The last hidden states (can be None, used in recurrent policies)
        :param episode_start: The last masks (can be None, used in recurrent policies)
            this correspond to beginning of episodes,
            where the hidden states of the RNN must be reset.
        :param deterministic: Whether or not to return deterministic actions.
        :return: the model's action and the next hidden state
            (used in recurrent policies)
        """
        # Switch to eval mode (this affects batch norm / dropout)
        self.set_training_mode(False)

        # Check for common mistake that the user does not mix Gym/VecEnv API
        # Tuple obs are not supported by SB3, so we can safely do that check
        if isinstance(observation, tuple) and len(observation) == 2 and isinstance(observation[1], dict):
            raise ValueError(
                "You have passed a tuple to the predict() function instead of a Numpy array or a Dict. "
                "You are probably mixing Gym API with SB3 VecEnv API: `obs, info = env.reset()` (Gym) "
                "vs `obs = vec_env.reset()` (SB3 VecEnv). "
            )
```

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

### TQC Model Setup

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

Sets up the TQC model, including creating aliases for actor and critic networks, initializing batch normalization statistics, and configuring the target entropy and entropy coefficient.

```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)
```

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

### OnPolicyAlgorithm._setup_model

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

Sets up the learning rate schedule, random seed, rollout buffer, and policy network. It also handles the selection of the appropriate rollout buffer class based on the observation space.

```APIDOC
## OnPolicyAlgorithm._setup_model

### Description
Sets up the learning rate schedule, random seed, rollout buffer, and policy network. It also handles the selection of the appropriate rollout buffer class based on the observation space.

### Method
`_setup_model(self) -> None`

### Parameters
None

### Returns
None
```

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

### Clone Repository and Navigate

Source: https://sb3-contrib.readthedocs.io/en/master/_sources/modules/ppo_mask.md.txt

Clone the repository for the experiment and navigate into the directory. This is the first step to replicate the results.

```bash
git clone git@github.com:kronion/microrts-ppo-comparison.git
cd microrts-ppo-comparison
```

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

### Setup Model - Stable-Baselines3

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

Initializes the learning rate schedule, sets the random seed, and configures the rollout buffer and policy network. It automatically selects the appropriate RolloutBuffer class based on the observation space.

```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()
```

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

### PPO Recurrent Model Setup

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

Sets up the recurrent PPO model, including the learning rate schedule, policy, and recurrent states. It selects the appropriate rollout buffer and initializes clipping schedules for policy and value functions.

```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)

```

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

### Get Action Distribution

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

Retrieves the action distribution for a given observation. It first extracts features, then passes them through the actor's MLP extractor to get latent representations, and finally uses these to form the distribution.

```python
def get_distribution(self, obs: PyTorchObs) -> Distribution:
        """
        Get the current policy distribution given the observations.

        :param obs:
        :return: the action distribution.
        """
        features = super().extract_features(obs, self.pi_features_extractor)
        latent_pi = self.mlp_extractor.forward_actor(features)
        return self._get_action_dist_from_latent(latent_pi)
```

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

### TQC Initialization Parameters

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

Defines the parameters for initializing the TQC agent, including learning rates, buffer sizes, and specific TQC parameters like target entropy and quantiles to drop.

```python
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,
```

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

### PPO Mask Model Setup

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

Sets up the PPO Mask model, including the learning rate schedule, random seed, policy, and rollout buffer. It enforces the use of MaskableActorCriticPolicy and selects the appropriate MaskableRolloutBuffer based on the observation space.

```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)

```

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

### Prediction

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

Method to get the policy action from an observation.

```APIDOC
## predict(_observation_, _state =None_, _episode_start =None_, _deterministic =False_, _action_masks =None_)

### Description
Get the policy action from an observation (and optional hidden state). Includes sugar-coating to handle different observations (e.g. normalizing images).

### Parameters
* **observation** (_ndarray_ | _dict_[_str_, _ndarray_]) – the input observation
* **state** (_tuple_[_ndarray_, _..._] | _None_) – The last hidden states (can be None, used in recurrent policies)
* **episode_start** (_ndarray_ | _None_) – The last masks (can be None, used in recurrent policies) this correspond to beginning of episodes, where the hidden states of the RNN must be reset.
* **deterministic** (_bool_) – Whether or not to return deterministic actions.
* **action_masks** (_ndarray_ | _None_)

### Returns
the model’s action and the next hidden state (used in recurrent policies)

### Return type
tuple[_ndarray_, tuple[_ndarray_, …] | None]
```

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

### 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
## Method: predict_values

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

### 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
the estimated values.

### Return type
Tensor
```

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

### Base Algorithm Initialization

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

Initializes the base algorithm with policy, device, and environment configurations. Handles policy class selection and environment wrapping. Includes checks for supported action spaces and observation types.

```python
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"

```

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

### _setup_lr_schedule

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

Transforms the learning rate into a callable schedule if it's not already.

```APIDOC
## _setup_lr_schedule

### Description
Transforms the learning rate into a callable schedule if it's not already.

### Method
`_setup_lr_schedule(self) -> None`
```

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

### QRDQNPolicy

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

The base policy class for QR-DQN. It handles the network building, optimizer setup, and forward pass.

```APIDOC
## QRDQNPolicy

### Description
Base policy class for QR-DQN.

### Methods
#### `__init__`
Initializes the QRDQNPolicy.

Parameters:
- `observation_space` (spaces.Space): The observation space of the environment.
- `action_space` (spaces.Discrete): The action space of the environment.
- `lr_schedule` (Schedule): The learning rate schedule.
- `n_quantiles` (int): The number of quantiles to use.
- `net_arch` (list[int] | None): The network architecture.
- `activation_fn` (type[nn.Module]): The activation function.
- `normalize_images` (bool): Whether to normalize images.
- `optimizer_class` (type[th.optim.Optimizer]): The optimizer class.
- `optimizer_kwargs` (dict[str, Any] | None): Additional optimizer keyword arguments.

#### `forward`
Forward pass of the policy.

Parameters:
- `obs` (PyTorchObs): The observation.
- `deterministic` (bool): Whether to use deterministic actions.

Returns:
- `th.Tensor`: The predicted action.

#### `set_training_mode`
Put the policy in either training or evaluation mode.

Parameters:
- `mode` (bool): If true, set to training mode, else set to evaluation mode.
```

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

### Clone and Run Benchmark

Source: https://sb3-contrib.readthedocs.io/en/master/_sources/modules/ppo_recurrent.md.txt

Instructions for cloning the rl-baselines3-zoo repository and running a benchmark for the PPO-LSTM algorithm against specified environments.

```bash
git clone https://github.com/DLR-RM/rl-baselines3-zoo
cd rl-baselines3-zoo

python train.py --algo ppo_lstm --env $ENV_ID --eval-episodes 10 --eval-freq 10000
```

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

### QR-DQN Initialization

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

Initializes the QR-DQN model with various hyperparameters. Note the specific optimizer epsilon adjustment for batch size.

```python
def __init__(
        self,
        policy: Type[QRPolicy],
        env: VecEnv | str | None,
        learning_rate: float = 1e-4,
        buffer_size: int = 1_000_000,
        learning_starts: int = 0,
        batch_size: int | None = None,
        tau: float = 1.0,
        gamma: float = 0.99,
        train_freq: int | tuple[int, str] = (1, "step"),
        gradient_steps: int = 1,
        exploration_fraction: float = 0.1,
        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()

```