RL Modules#
With the RLModule class in RLlib’s new API stack, you write custom models, including complex multi-network setups often found in multi-agent or model-based algorithms.
RLModule is the main neural network class and exposes three public methods, each corresponding to a distinct phase in the reinforcement learning cycle:
forward_exploration()computes actions during data collection when RLlib uses the data for a subsequent training step, balancing exploration and exploitation.forward_inference()computes actions for evaluation and production, which often need to be greedy or less stochastic.forward_train()manages the training phase, performing calculations required to compute losses, such as Q-values in a DQN model, value function predictions in a PG-style setup, or world-model predictions in model-based algorithms.
RLModule overview: (left) A plain RLModule contains the
neural network RLlib uses for computations,
for example, a policy network written in PyTorch, and exposes the three forward methods:
forward_exploration() for sample collection,
forward_inference() for production/deployment, and
forward_train() for computing loss function inputs when training.
(right) A MultiRLModule may contain one
or more sub-RLModules, each identified by a ModuleID, so you can implement
arbitrarily complex multi-network or multi-agent architectures and algorithms.#
Enable the RLModule API in the AlgorithmConfig#
In the new API stack, activated by default, RLlib exclusively uses RLModules.
If you’re working with a legacy config or want to migrate ModelV2 or Policy classes to the new API stack, see the new API stack migration guide.
If you configured the Algorithm to the old API stack, use the api_stack() method to switch:
from ray.rllib.algorithms.algorithm_config import AlgorithmConfig
config = (
AlgorithmConfig()
.api_stack(
enable_rl_module_and_learner=True,
enable_env_runner_and_connector_v2=True,
)
)
Default RLModules#
If you don’t specify module-related settings in the AlgorithmConfig, RLlib uses the respective algorithm’s default RLModule, which is an appropriate choice for initial experimentation and benchmarking. All default RLModules support 1D-tensor and image observations of the form [width] x [height] x [channels].
Note
For discrete or more complex input observation spaces such as dictionaries, use the FlattenObservations connector piece as follows:
from ray.rllib.algorithms.ppo import PPOConfig
from ray.rllib.connectors.env_to_module import FlattenObservations
config = (
PPOConfig()
# FrozenLake has a discrete observation space (ints).
.environment("FrozenLake-v1")
# `FlattenObservations` converts int observations to one-hot.
.env_runners(env_to_module_connector=lambda env: FlattenObservations())
)
All default models also offer configurable architecture choices. You can set the number and size of the layers, either Dense or Conv2D, their activations and initializations, and the automatic LSTM-wrapping behavior.
Use the DefaultModelConfig dataclass to configure any default model in RLlib. Use this class only for default models. When writing your own custom RLModules, use plain Python dicts to define the model configurations. For how to write and configure your custom RLModules, see Implementing custom RLModules.
Configure default MLP nets#
To train a simple multi-layer perceptron (MLP) policy, which only contains dense layers, with PPO and the default RLModule, configure your experiment as follows:
from ray.rllib.algorithms.ppo import PPOConfig
from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
config = (
PPOConfig()
.environment("CartPole-v1")
.rl_module(
# Use a non-default 32,32-stack with ReLU activations.
model_config=DefaultModelConfig(
fcnet_hiddens=[32, 32],
fcnet_activation="relu",
)
)
)
The following is the complete list of all supported fcnet_.. options:
#: List containing the sizes (number of nodes) of a fully connected (MLP) stack.
#: Note that in an encoder-based default architecture with a policy head (and
#: possible value head), this setting only affects the encoder component. To set the
#: policy (and value) head sizes, use `head_fcnet_hiddens`, instead. For example,
#: if you set `fcnet_hiddens=[32, 32]` and `head_fcnet_hiddens=[64]`, you would get
#: an RLModule with a [32, 32] encoder, a [64, act-dim] policy head, and a [64, 1]
#: value head (if applicable).
fcnet_hiddens: List[int] = field(default_factory=lambda: [256, 256])
#: Activation function descriptor for the stack configured by `fcnet_hiddens`.
#: Supported values are: 'tanh', 'relu', 'swish' (or 'silu', which is the same),
#: and 'linear' (or None).
fcnet_activation: str = "tanh"
#: Initializer function or class descriptor for the weight/kernel matrices in the
#: stack configured by `fcnet_hiddens`. Supported values are the initializer names
#: (str), classes or functions listed by the frameworks (`torch`). See
#: https://pytorch.org/docs/stable/nn.init.html for `torch`. If `None` (default),
#: the default initializer defined by `torch` is used.
fcnet_kernel_initializer: Optional[Union[str, Callable]] = None
#: Kwargs passed into the initializer function defined through
#: `fcnet_kernel_initializer`.
fcnet_kernel_initializer_kwargs: Optional[dict] = None
#: Initializer function or class descriptor for the bias vectors in the stack
#: configured by `fcnet_hiddens`. Supported values are the initializer names (str),
#: classes or functions listed by the frameworks (`torch`). See
#: https://pytorch.org/docs/stable/nn.init.html for `torch`. If `None` (default),
#: the default initializer defined by `torch` is used.
fcnet_bias_initializer: Optional[Union[str, Callable]] = None
#: Kwargs passed into the initializer function defined through
#: `fcnet_bias_initializer`.
fcnet_bias_initializer_kwargs: Optional[dict] = None
#: Whether to insert LayerNorm after each hidden layer in the encoder stack
#: configured by `fcnet_hiddens`.
fcnet_use_layernorm: bool = False
Configure default CNN nets#
For image-based environments such as Atari, use the conv_.. fields in DefaultModelConfig to configure the convolutional neural network (CNN) stack.
You might have to check whether your CNN configuration works with the incoming observation image dimensions. For example, for an Atari environment, you can use RLlib’s Atari wrapper utility, which performs resizing (default 64x64) and gray scaling (default True), frame stacking (default None), frame skipping (default 4), normalization (from uint8 to float32), and applies up to 30 “noop” actions after a reset, which aren’t part of the episode:
import gymnasium as gym # `pip install gymnasium[atari,accept-rom-license]`
from ray.rllib.algorithms.ppo import PPOConfig
from ray.rllib.env.wrappers.atari_wrappers import wrap_atari_for_new_api_stack
from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
from ray.tune import register_env
register_env(
"image_env",
lambda _: wrap_atari_for_new_api_stack(
gym.make("ale_py:ALE/Pong-v5"),
dim=64, # resize original observation to 64x64x3
framestack=4,
)
)
config = (
PPOConfig()
.environment("image_env")
.rl_module(
model_config=DefaultModelConfig(
# Use a DreamerV3-style CNN stack for 64x64 images.
conv_filters=[
[16, 4, 2], # 1st CNN layer: num_filters, kernel, stride(, padding)?
[32, 4, 2], # 2nd CNN layer
[64, 4, 2], # etc..
[128, 4, 2],
],
conv_activation="silu",
# After the last CNN, the default model flattens, then adds an optional MLP.
head_fcnet_hiddens=[256],
)
)
)
Configure LSTM#
To auto-wrap your default encoder with an extra LSTM layer so your model can learn in non-Markovian, partially observable environments, use the DefaultModelConfig.use_lstm setting together with the DefaultModelConfig.lstm_cell_size and DefaultModelConfig.max_seq_len settings. For a tuned example, see an example that uses a default RLModule with an LSTM layer.
Construct RLModule instances#
RLlib offers a standardized approach for constructing RLModule instances for both single-module and multi-module use cases. An example of a single-module use case is a single-agent experiment. Examples of multi-module use cases are multi-agent learning or other multi-NN setups.
Construction through the class constructor#
The most direct way to construct your RLModule is through its constructor:
import gymnasium as gym
from ray.rllib.algorithms.bc.torch.default_bc_torch_rl_module import DefaultBCTorchRLModule
# Create an env object to know the spaces.
env = gym.make("CartPole-v1")
# Construct the actual RLModule object.
rl_module = DefaultBCTorchRLModule(
observation_space=env.observation_space,
action_space=env.action_space,
# A custom dict that's accessible inside your class as `self.model_config`.
model_config={"fcnet_hiddens": [64]},
)
Note
If you have a checkpoint of an Algorithm or an individual RLModule, see Creating instances with from_checkpoint for how to recreate your RLModule from disk.
Construction through RLModuleSpecs#
Because RLlib is a distributed RL library and needs to create more than one copy of your RLModule, you can use RLModuleSpec objects to define how RLlib should construct each copy during the algorithm’s setup process. The algorithm passes the spec to all subcomponents that need a copy of your RLModule.
Creating an RLModuleSpec is analogous to the RLModule constructor:
import gymnasium as gym
from ray.rllib.algorithms.bc.torch.default_bc_torch_rl_module import DefaultBCTorchRLModule
from ray.rllib.core.rl_module.rl_module import RLModuleSpec
# Create an env object to know the spaces.
env = gym.make("CartPole-v1")
# First construct the spec.
spec = RLModuleSpec(
module_class=DefaultBCTorchRLModule,
observation_space=env.observation_space,
action_space=env.action_space,
# A custom dict that's accessible inside your class as `self.model_config`.
model_config={"fcnet_hiddens": [64]},
)
# Then, build the RLModule through the spec's `build()` method.
rl_module = spec.build()
import gymnasium as gym
from ray.rllib.algorithms.bc.torch.default_bc_torch_rl_module import DefaultBCTorchRLModule
from ray.rllib.core.rl_module.rl_module import RLModuleSpec
from ray.rllib.core.rl_module.multi_rl_module import MultiRLModuleSpec
# First construct the MultiRLModuleSpec.
spec = MultiRLModuleSpec(
rl_module_specs={
"module_1": RLModuleSpec(
module_class=DefaultBCTorchRLModule,
# Define the spaces for only this sub-module.
observation_space=gym.spaces.Box(low=-1, high=1, shape=(10,)),
action_space=gym.spaces.Discrete(2),
# A custom dict that's accessible inside your class as
# `self.model_config`.
model_config={"fcnet_hiddens": [32]},
),
"module_2": RLModuleSpec(
module_class=DefaultBCTorchRLModule,
# Define the spaces for only this sub-module.
observation_space=gym.spaces.Box(low=-1, high=1, shape=(5,)),
action_space=gym.spaces.Discrete(2),
# A custom dict that's accessible inside your class as
# `self.model_config`.
model_config={"fcnet_hiddens": [16]},
),
},
)
# Construct the actual MultiRLModule instance with .build():
multi_rl_module = spec.build()
You can pass the RLModuleSpec instances to your AlgorithmConfig to tell RLlib to use the particular module class and constructor arguments:
from ray.rllib.algorithms.ppo import PPOConfig
from ray.rllib.core.rl_module.rl_module import RLModuleSpec
config = (
PPOConfig()
.environment("CartPole-v1")
.rl_module(
rl_module_spec=RLModuleSpec(
module_class=MyRLModuleClass,
model_config={"some_key": "some_setting"},
),
)
)
ppo = config.build()
print(ppo.get_module())
Note
Often when creating an RLModuleSpec, you don’t have to define attributes such as observation_space or action_space because RLlib automatically infers these attributes from the environment or other configuration parameters.
from ray.rllib.algorithms.ppo import PPOConfig
from ray.rllib.core.rl_module.rl_module import RLModuleSpec
from ray.rllib.core.rl_module.multi_rl_module import MultiRLModuleSpec
from ray.rllib.examples.envs.classes.multi_agent import MultiAgentCartPole
config = (
PPOConfig()
.environment(MultiAgentCartPole, env_config={"num_agents": 2})
.multi_agent(
# Both agents (0 and 1) map to the same policy, so they share
# a single RLModule.
policies={"p0"},
policy_mapping_fn=lambda agent_id, episode, **kw: "p0",
)
.rl_module(
rl_module_spec=MultiRLModuleSpec(
rl_module_specs={
"p0": RLModuleSpec(
module_class=MyRLModuleClass,
model_config={"some_key": "some_setting"},
),
},
),
)
)
ppo = config.build()
print(ppo.get_module())
from ray.rllib.algorithms.ppo import PPOConfig
from ray.rllib.core.rl_module.rl_module import RLModuleSpec
from ray.rllib.core.rl_module.multi_rl_module import MultiRLModuleSpec
from ray.rllib.examples.envs.classes.multi_agent import MultiAgentCartPole
config = (
PPOConfig()
.environment(MultiAgentCartPole, env_config={"num_agents": 2})
.multi_agent(
policies={"p0", "p1"},
# Agent IDs of `MultiAgentCartPole` are 0 and 1, mapping to
# "p0" and "p1", respectively.
policy_mapping_fn=lambda agent_id, episode, **kw: f"p{agent_id}"
)
.rl_module(
rl_module_spec=MultiRLModuleSpec(
# Agents (0 and 1) use different (single) RLModules.
rl_module_specs={
"p0": RLModuleSpec(
module_class=MyRLModuleClass,
# Small network.
model_config={"fcnet_hiddens": [32, 32]},
),
"p1": RLModuleSpec(
module_class=MyRLModuleClass,
# Large network.
model_config={"fcnet_hiddens": [128, 128]},
),
},
),
)
)
ppo = config.build()
print(ppo.get_module())
Implement custom RLModules#
To implement your own neural network architecture and computation logic, subclass TorchRLModule for any single-agent learning experiment or for independent multi-agent learning.
For more advanced multi-agent use cases such as ones with shared communication between agents, or any multi-model use cases, subclass the MultiRLModule class instead.
Note
An alternative to subclassing TorchRLModule is to directly subclass your Algorithm’s default RLModule. For example, to use PPO, subclass DefaultPPOTorchRLModule. In this case, carefully study the existing default model to understand how to override the setup(), the _forward_() methods, and possibly some algo-specific API methods. See Algorithm-specific RLModule APIs for how to determine which APIs your algorithm requires you to implement.
The setup() method#
First implement the setup() method, where you add the NN subcomponents you need and assign them to class attributes of your choice.
Call super().setup() in your implementation.
You also have access to the following attributes anywhere in the class, including in setup():
self.observation_spaceself.action_spaceself.inference_onlyself.model_config, a dict with any custom config settings
import torch
from ray.rllib.core.rl_module.torch.torch_rl_module import TorchRLModule
class MyTorchPolicy(TorchRLModule):
def setup(self):
# You have access here to the following already set attributes:
# self.observation_space
# self.action_space
# self.inference_only
# self.model_config # <- a dict with custom settings
# Use the observation space (if a Box) to infer the input dimension.
input_dim = self.observation_space.shape[0]
# Use the model_config dict to extract the hidden dimension.
hidden_dim = self.model_config["fcnet_hiddens"][0]
# Use the action space to infer the number of output nodes.
output_dim = self.action_space.n
# Build all the layers and subcomponents here you need for the
# RLModule's forward passes.
self._pi_head = torch.nn.Sequential(
torch.nn.Linear(input_dim, hidden_dim),
torch.nn.ReLU(),
torch.nn.Linear(hidden_dim, output_dim),
)
Forward methods#
To implement the forward computation logic, you have two options. Either define a generic forward behavior by overriding the private _forward() method, which RLlib then uses throughout the model’s lifecycle, or, for more granularity, define the following three private methods:
_forward_exploration(): Forward pass for computing exploration actions for collecting training data._forward_inference(): Forward pass for action inference, such as greedy._forward_train(): Forward pass for computing loss function inputs for a training update.
For custom _forward(), _forward_inference(), and _forward_exploration() methods, you must return a dictionary that contains the key actions, the key action_dist_inputs, or both.
If you return the actions key from your forward method:
RLlib uses the provided actions as-is.
If you also return the
action_dist_inputskey, RLlib creates aDistributioninstance from the parameters under that key. Forforward_exploration(), RLlib also computes action probabilities and log probabilities for the given actions automatically. See Custom action distributions.
If you don’t return the actions key from your forward method:
You must return the
action_dist_inputskey from your_forward_exploration()and_forward_inference()methods.RLlib creates a
Distributioninstance from the parameters under that key and samples actions from that distribution. See Custom action distributions.For
_forward_exploration(), RLlib also computes action probability and log probability values from the sampled actions automatically.
Note
For _forward_inference(), RLlib always makes the generated distributions from returned key action_dist_inputs deterministic first through the to_deterministic() utility before a possible action sample step. For example, RLlib reduces the sampling from a Categorical distribution to selecting the argmax actions from the distribution logits or probabilities. If you return the “actions” key, RLlib skips that sampling step.
from ray.rllib.core import Columns, TorchRLModule
class MyTorchPolicy(TorchRLModule):
...
def _forward_inference(self, batch):
...
return {
Columns.ACTIONS: ... # RLlib uses these actions as-is
}
def _forward_exploration(self, batch):
...
return {
Columns.ACTIONS: ..., # RLlib uses these actions as-is (no sampling step!)
Columns.ACTION_DIST_INPUTS: ... # If provided, RLlib uses these dist inputs to compute probs and logp.
}
from ray.rllib.core import Columns, TorchRLModule
class MyTorchPolicy(TorchRLModule):
...
def _forward_inference(self, batch):
...
return {
# RLlib:
# - Generates distribution from ACTION_DIST_INPUTS parameters.
# - Converts distribution to a deterministic equivalent.
# - Samples from the deterministic distribution.
Columns.ACTION_DIST_INPUTS: ...
}
def _forward_exploration(self, batch):
...
return {
# RLlib:
# - Generates distribution from ACTION_DIST_INPUTS parameters.
# - Samples from the stochastic distribution.
# - Computes action probs and logs automatically using the sampled
# actions and the distribution.
Columns.ACTION_DIST_INPUTS: ...
}
Never override the constructor, __init__. The RLModule class’s constructor requires the following arguments, and it receives them properly when you call a spec’s build() method:
observation_space: The observation space after passing through all connectors. This is the actual input space for the model after all preprocessing steps.action_space: The action space of the environment.inference_only: Whether RLlib should build the RLModule in inference-only mode, dropping subcomponents that it only needs for learning.model_config: The model config, which is either a custom dictionary for custom RLModules or aDefaultModelConfigdataclass object, which is only for RLlib’s default models. Define model hyper-parameters such as the number of layers and the type of activation in this object.
Algorithm-specific RLModule APIs#
The algorithm you choose to use with your RLModule affects the structure of the final custom module to some extent. Each Algorithm class has a fixed set of APIs that all RLModules trained by that algorithm need to implement.
To find out what APIs your Algorithms require, do the following:
# Import the config of the algorithm of your choice.
from ray.rllib.algorithms.sac import SACConfig
# Print out the abstract APIs, you need to subclass from and whose
# abstract methods you need to implement, besides the ``setup()`` and ``_forward_..()``
# methods.
print(
SACConfig()
.get_default_learner_class()
.rl_module_required_apis()
)
Note
You didn’t implement any APIs in the preceding example module, because you hadn’t considered training it with any particular algorithm yet. You can find examples of custom RLModule classes implementing the SelfSupervisedLossAPI and thus ready to train with PPO in the tiny_atari_cnn_rlm example and in the lstm_containing_rlm example.
You can mix supervised losses into any RLlib algorithm through the SelfSupervisedLossAPI. Your Learner actors automatically call the implemented compute_self_supervised_loss() method to compute the model’s own loss passing it the outputs of the forward_train() call.
See the example script that uses a self-supervised loss RLModule. You can define losses over either policy evaluation inputs or data read from offline storage. Set the learner_only attribute to True in your custom RLModuleSpec if you don’t need the self-supervised model for collecting samples in your EnvRunner actors. In this case, you might also need an extra Learner connector piece to make sure your RLModule receives data to learn.
End-to-end example#
The following working end-to-end example puts together the elements of the custom RLModule you implemented:
import torch
from ray.rllib.core.columns import Columns
from ray.rllib.core.rl_module.torch import TorchRLModule
class VPGTorchRLModule(TorchRLModule):
"""A simple VPG (vanilla policy gradient)-style RLModule for testing purposes.
Use this as a minimum, bare-bones example implementation of a custom TorchRLModule.
"""
def setup(self):
"""Use this method to create all the model components that you require.
Feel free to access the following useful properties in this class:
- `self.model_config`: The config dict for this RLModule class,
which should contain flexible settings, for example: {"hiddens": [256, 256]}.
- `self.observation|action_space`: The observation and action space that
this RLModule is subject to. Note that the observation space might not be the
exact space from your env, but that it might have already gone through
preprocessing through a connector pipeline (for example, flattening,
frame-stacking, mean/std-filtering, etc..).
- `self.inference_only`: If True, this model should be built only for inference
purposes, in which case you may want to exclude any components that are not used
for computing actions, for example a value function branch.
"""
input_dim = self.observation_space.shape[0]
hidden_dim = self.model_config["hidden_dim"]
output_dim = self.action_space.n
self._policy_net = torch.nn.Sequential(
torch.nn.Linear(input_dim, hidden_dim),
torch.nn.ReLU(),
torch.nn.Linear(hidden_dim, output_dim),
)
def _forward(self, batch, **kwargs):
# Push the observations from the batch through our `self._policy_net`.
action_logits = self._policy_net(batch[Columns.OBS])
# Return parameters for the (default) action distribution, which is
# `TorchCategorical` (due to our action space being `gym.spaces.Discrete`).
return {Columns.ACTION_DIST_INPUTS: action_logits}
# If you need more granularity between the different forward behaviors during
# the different phases of the module's lifecycle, implement three different
# forward methods. Thereby, it is recommended to put the inference and
# exploration versions inside a `with torch.no_grad()` context for better
# performance.
# def _forward_train(self, batch):
# ...
#
# def _forward_inference(self, batch):
# with torch.no_grad():
# return self._forward_train(batch)
#
# def _forward_exploration(self, batch):
# with torch.no_grad():
# return self._forward_train(batch)
Custom action distributions#
The preceding examples rely on RLModule using the correct action distribution with the computed ACTION_DIST_INPUTS returned by the forward methods. RLlib picks a default distribution class based on the action space, which is TorchCategorical for Discrete action spaces and TorchDiagGaussian for Box action spaces.
To use a different distribution class and return parameters for this distribution’s constructor from your RLModule forward methods, set the action_dist_cls attribute inside the setup() method of your RLModule.
See the example script that introduces a temperature parameter on top of a Categorical distribution.
To specify different distribution classes for the different forward methods of your RLModule, override the following methods and return a different distribution class from each:
Note
If you only return ACTION_DIST_INPUTS from your forward methods, RLlib automatically uses the to_deterministic() method of the distribution returned by your get_inference_action_dist_cls().
See torch_distributions.py for common distribution implementations.
Auto-regressive action distributions#
In an action space with multiple components, for example Tuple(a1, a2), you might want to condition the sampling of a2 on the sampled value of a1, such that a2_sampled ~ P(a2 | a1_sampled, obs). In the default, non-autoregressive case, RLlib uses a default model with an independent TorchMultiDistribution and samples a1 and a2 independently. This makes learning impossible in environments where the sampling of one action component must depend on another, already-sampled component. See an example of a “correlated actions” environment.
To write a custom RLModule that samples the action components as previously described, carefully implement its forward logic.
Find an example of such an autoregressive action model.
You implement the main action sampling logic in the _forward_...() methods:
def _pi(self, obs, inference: bool):
# Prior forward pass and sample a1.
prior_out = self._prior_net(obs)
dist_a1 = TorchCategorical.from_logits(prior_out)
if inference:
dist_a1 = dist_a1.to_deterministic()
a1 = dist_a1.sample()
# Posterior forward pass and sample a2.
posterior_batch = torch.cat(
[obs, one_hot(a1, self.action_space[0])],
dim=-1,
)
posterior_out = self._posterior_net(posterior_batch)
dist_a2 = TorchDiagGaussian.from_logits(posterior_out)
if inference:
dist_a2 = dist_a2.to_deterministic()
a2 = dist_a2.sample()
actions = (a1, a2)
# We need logp and distribution parameters for the loss.
return {
Columns.ACTION_LOGP: (
TorchMultiDistribution((dist_a1, dist_a2)).logp(actions)
),
Columns.ACTION_DIST_INPUTS: torch.cat([prior_out, posterior_out], dim=-1),
Columns.ACTIONS: actions,
}
Implement custom MultiRLModules#
For multi-module setups, RLlib provides the MultiRLModule class, whose default implementation is a dictionary of individual RLModule objects, one for each submodule and identified by a ModuleID.
The base-class MultiRLModule implementation works for most use cases that need independent neural networks. For a complex multi-network or multi-agent use case where agents share one or more neural networks, inherit from this class and override the default implementation.
The following code snippets create a custom multi-agent RLModule with two “policy head” modules, which share the same encoder, the third network in the MultiRLModule. The encoder receives the raw observations from the environment and outputs embedding vectors that then serve as input for the two policy heads to compute the agents’ actions.
class VPGMultiRLModuleWithSharedEncoder(MultiRLModule):
"""VPG (vanilla pol. gradient)-style MultiRLModule handling a shared encoder.
def setup(self):
# Call the super's setup().
super().setup()
# Assert, we have the shared encoder submodule.
assert SHARED_ENCODER_ID in self._rl_modules and len(self._rl_modules) > 1
# Assign the encoder to a convenience attribute.
self.encoder = self._rl_modules[SHARED_ENCODER_ID]
def _forward(self, batch, forward_type, **kwargs):
# Collect our policies' outputs in this dict.
fwd_out = {}
# Loop through the policy nets (through the given batch's keys).
for policy_id, policy_batch in batch.items():
# Feed this policy's observation into the shared encoder
encoder_output = self.encoder._forward(batch[policy_id])
policy_batch[ENCODER_OUT] = encoder_output[ENCODER_OUT]
# Get the desired module
m = getattr(self._rl_modules[policy_id], forward_type)
# Pass the policy's embeddings through the policy net.
fwd_out[policy_id] = m(batch[policy_id], **kwargs)
return fwd_out
# These methods could probably stand to be adjusted in MultiRLModule using something like this, so that subclasses that tweak _forward don't need to rewrite all of them. The prior implementation errored out because of this issue.
@override(MultiRLModule)
def _forward_inference(
self, batch: Dict[str, Any], **kwargs
) -> Union[Dict[str, Any], Dict[ModuleID, Dict[str, Any]]]:
return self._forward(batch, "_forward_inference", **kwargs)
@override(MultiRLModule)
def _forward_exploration(
self, batch: Dict[str, Any], **kwargs
) -> Union[Dict[str, Any], Dict[ModuleID, Dict[str, Any]]]:
return self._forward(batch, "_forward_exploration", **kwargs)
@override(MultiRLModule)
def _forward_train(
self, batch: Dict[str, Any], **kwargs
) -> Union[Dict[str, Any], Dict[ModuleID, Dict[str, Any]]]:
return self._forward(batch, "_forward_train", **kwargs)
Within the MultiRLModule, you need two policy sub-RLModules. They can be of the same class, which you implement as follows:
class VPGPolicyAfterSharedEncoder(TorchRLModule):
"""A VPG (vanilla pol. gradient)-style RLModule using a shared encoder.
def setup(self):
super().setup()
# Incoming feature dim from the shared encoder.
embedding_dim = self.model_config["embedding_dim"]
hidden_dim = self.model_config["hidden_dim"]
self._pi_head = torch.nn.Sequential(
torch.nn.Linear(embedding_dim, hidden_dim),
torch.nn.ReLU(),
torch.nn.Linear(hidden_dim, self.action_space.n),
)
def _forward(self, batch, **kwargs):
embeddings = batch[ENCODER_OUT] # Get the output of the encoder
logits = self._pi_head(embeddings)
return {Columns.ACTION_DIST_INPUTS: logits}
Finally, the shared encoder RLModule should look similar to this:
class SharedEncoder(TorchRLModule):
"""A shared encoder that can be used with `VPGMultiRLModuleWithSharedEncoder`."""
def setup(self):
super().setup()
input_dim = self.observation_space.shape[0]
embedding_dim = self.model_config["embedding_dim"]
# A very simple encoder network.
self._net = torch.nn.Sequential(
torch.nn.Linear(input_dim, embedding_dim),
)
def _forward(self, batch, **kwargs):
# Pass observations through the net and return outputs.
return {ENCODER_OUT: self._net(batch[Columns.OBS])}
To plug the custom MultiRLModule from the first tab into your algorithm’s config, create a MultiRLModuleSpec with the new class and its constructor settings. Also create one RLModuleSpec for each agent and for the shared encoder RLModule, because RLlib requires their observation and action spaces and their model hyper-parameters:
import gymnasium as gym
from ray.rllib.core.rl_module.rl_module import RLModuleSpec
from ray.rllib.core.rl_module.multi_rl_module import MultiRLModuleSpec
from ray.rllib.examples.algorithms.classes.vpg import VPGConfig
from ray.rllib.examples.learners.classes.vpg_torch_learner_shared_optimizer import VPGTorchLearnerSharedOptimizer
from ray.rllib.examples.envs.classes.multi_agent import MultiAgentCartPole
from ray.rllib.examples.rl_modules.classes.vpg_using_shared_encoder_rlm import (
SHARED_ENCODER_ID,
SharedEncoder,
VPGPolicyAfterSharedEncoder,
VPGMultiRLModuleWithSharedEncoder,
)
single_agent_env = gym.make("CartPole-v1")
EMBEDDING_DIM = 64 # encoder output dim
config = (
VPGConfig()
.environment(MultiAgentCartPole, env_config={"num_agents": 2})
.training(
learner_class=VPGTorchLearnerSharedOptimizer,
)
.multi_agent(
# Declare the two policies trained.
policies={"p0", "p1"},
# Agent IDs of `MultiAgentCartPole` are 0 and 1. They are mapped to
# the two policies with ModuleIDs "p0" and "p1", respectively.
policy_mapping_fn=lambda agent_id, episode, **kw: f"p{agent_id}"
)
.rl_module(
rl_module_spec=MultiRLModuleSpec(
multi_rl_module_class=VPGMultiRLModuleWithSharedEncoder,
rl_module_specs={
# Shared encoder.
SHARED_ENCODER_ID: RLModuleSpec(
module_class=SharedEncoder,
model_config={"embedding_dim": EMBEDDING_DIM},
observation_space=single_agent_env.observation_space,
action_space=single_agent_env.action_space,
),
# Large policy net.
"p0": RLModuleSpec(
module_class=VPGPolicyAfterSharedEncoder,
model_config={
"embedding_dim": EMBEDDING_DIM,
"hidden_dim": 1024,
},
),
# Small policy net.
"p1": RLModuleSpec(
module_class=VPGPolicyAfterSharedEncoder,
model_config={
"embedding_dim": EMBEDDING_DIM,
"hidden_dim": 64,
},
),
},
),
)
)
algo = config.build_algo()
print(algo.train())
Note
To properly learn with the preceding setup, write and use a specific multi-agent Learner that can handle the shared encoder. This Learner should have only a single optimizer that updates all three submodules, the encoder and the two policy nets, to stabilize learning. With the standard “one-optimizer-per-module” Learners, the two optimizers for policy 1 and policy 2 take turns updating the same shared encoder, which leads to learning instabilities.
Checkpoint RLModules#
You can checkpoint RLModule instances with their save_to_path() method. If you already have an instantiated RLModule and want to load new model weights into it from an existing checkpoint, use the restore_from_path() method.
The following examples show how to use these methods outside of an RLlib Algorithm or together with one.
Create an RLModule checkpoint#
import tempfile
import gymnasium as gym
from ray.rllib.algorithms.ppo.torch.default_ppo_torch_rl_module import DefaultPPOTorchRLModule
from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
env = gym.make("CartPole-v1")
# Create an RLModule to later checkpoint.
rl_module = DefaultPPOTorchRLModule(
observation_space=env.observation_space,
action_space=env.action_space,
model_config=DefaultModelConfig(fcnet_hiddens=[32]),
)
# Finally, write the RLModule checkpoint.
module_ckpt_path = tempfile.mkdtemp()
rl_module.save_to_path(module_ckpt_path)
Create an RLModule from a checkpoint#
If you have an RLModule checkpoint saved and want to create a new RLModule directly from it, use the from_checkpoint() method:
from ray.rllib.core.rl_module.rl_module import RLModule
# Create a new RLModule from the checkpoint.
new_module = RLModule.from_checkpoint(module_ckpt_path)
Load an RLModule checkpoint into a running Algorithm#
from ray.rllib.algorithms.ppo import PPOConfig
# Create a new Algorithm (with the changed module config: 32 units instead of the
# default 256; otherwise loading the state of ``module`` fails due to a shape
# mismatch).
config = (
PPOConfig()
.environment("CartPole-v1")
.rl_module(model_config=DefaultModelConfig(fcnet_hiddens=[32]))
)
ppo = config.build()
You can load the saved RLModule state from the preceding module.save_to_path() directly into the running Algorithm’s RLModules. This updates all RLModules within the algorithm, both those in the Learner workers and those in the EnvRunners.
ppo.restore_from_path(
module_ckpt_path, # <- NOT an Algorithm checkpoint, but single-agent RLModule one.
# Therefore, we have to provide the exact path (of RLlib components) down
# to the individual RLModule within the algorithm, which is:
component="learner_group/learner/rl_module/default_policy",
)