Stable Baselines3 Contrib

CrossQ is an algorithm that utilizes batch normalization in the critic network and removes target networks to improve sample efficiency in off-policy deep reinforcement learning. It is designed to work without requiring high update-to-data ratios.

Key Constraints:

• Supported Action Spaces: Box (Continuous).
• Unsupported Action Spaces: Discrete, MultiDiscrete, MultiBinary.
• Supported Observation Spaces: Box.
• Unsupported Observation Spaces: Dict.
• Policy Support: Currently supports MlpPolicy. CnnPolicy (for images) is not yet implemented.
• Recurrent Policies: Not supported.
• Multi-processing: Supported.

Architecture Note: The default network architecture for the q-value function is [1024, 1024], which is a compromise between speed and performance compared to the original paper's [2048, 2048].

from sb3_contrib import CrossQ

# Initialize the model with MlpPolicy for a continuous control task
model = CrossQ("MlpPolicy", "Walker2d-v4")

# Train the agent
model.learn(total_timesteps=1_000_000)

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

QR-DQN is a distributional reinforcement learning algorithm that models the distribution over returns using quantile regression, rather than just predicting the mean return like standard DQN.

Supported Policies

• MlpPolicy: For vector-based observations.
• CnnPolicy: For image-based observations.
• MultiInputPolicy: For observations containing both vector and image data.

Compatibility

• Action Spaces: Supports Discrete actions.
• Observation Spaces: Supports Discrete, Box, MultiDiscrete, MultiBinary, and Dict spaces.
• Recurrent Policies: Not supported (❌).
• Multi-processing: Supported (✔️).

import gymnasium as gym
from sb3_contrib import QRDQN

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

# Configure quantiles via policy_kwargs
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")

# Loading and using the model
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()

Maskable PPO is an implementation of Proximal Policy Optimization (PPO) that supports invalid action masking. This allows the agent to ignore actions that are invalid in the current state, improving learning efficiency.

Key Behaviors:

• It behaves identically to SB3's standard PPO unless the environment is wrapped with ActionMasker.
• If an ActionMasker is detected, masks are automatically retrieved and used during training.
• Supported Spaces:
  • Discrete: Action (✔️), Observation (✔️)
  • MultiDiscrete: Action (✔️), Observation (✔️)
  • MultiBinary: Action (✔️), Observation (✔️)
  • Box: Action (❌), Observation (✔️)
  • Dict: Action (❌), Observation (✔️)

Important Requirements:

• Recurrent policies are not supported (❌).
• Multi-processing is supported (✔️).
• Evaluation: You MUST use MaskableEvalCallback from sb3_contrib.common.maskable.callbacks instead of the standard EvalCallback.
• Policy Evaluation: You MUST use evaluate_policy from sb3_contrib.common.maskable.evaluation instead of the standard SB3 version.
• SubprocVecEnv: If using SubprocVecEnv, you cannot use the ActionMasker wrapper; you must implement the action_masks method directly inside the environment.

from sb3_contrib import MaskablePPO

# Standard usage with an environment that has an action mask
model = MaskablePPO("MlpPolicy", env, verbose=1)
model.learn(total_timesteps=10_000)

TQC is an off-policy actor-critic algorithm designed to control overestimation bias by using quantile regression to predict a distribution for the value function instead of a single mean value. It builds upon SAC, TD3, and QR-DQN and truncates predicted quantiles to improve stability.

Supported Spaces

Supported Policies

• MlpPolicy
• CnnPolicy
• MultiInputPolicy

import gymnasium as gym
from sb3_contrib import TQC

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

# Configure policy arguments like number of critics and quantiles
policy_kwargs = dict(n_critics=2, n_quantiles=25)

# Initialize TQC
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 model
model.save("tqc_pendulum")

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

# Use the model for prediction
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()

If your custom environment uses a different method name than the standard for returning masks, you can use the ActionMasker wrapper to bridge it to MaskablePPO.

To use this, define a mask_fn that takes the environment and returns a numpy.ndarray representing the valid actions (where True indicates a valid action and False indicates an invalid one).

import gymnasium as gym
import numpy as np
from sb3_contrib.common.wrappers import ActionMasker
from sb3_contrib.ppo_mask import MaskablePPO
from sb3_contrib.common.maskable.policies import MaskableActorCriticPolicy

def mask_fn(env: gym.Env) -> np.ndarray:
    # Return the mask from your environment's specific method
    return env.valid_action_mask()

env = ...  # Your custom environment
env = ActionMasker(env, mask_fn)

# MaskablePPO will now automatically detect the wrapper and use the masks
model = MaskablePPO(MaskableActorCriticPolicy, env, verbose=1)
model.learn()