from typing import Any, Dict, Optional from ray.rllib.core.columns import Columns from ray.rllib.core.learner.utils import make_target_network from ray.rllib.core.rl_module.apis import ( TARGET_NETWORK_ACTION_DIST_INPUTS, TargetNetworkAPI, ValueFunctionAPI, ) from ray.rllib.core.rl_module.torch import TorchRLModule from ray.rllib.models.torch.misc import ( normc_initializer, same_padding, valid_padding, ) from ray.rllib.utils.annotations import override from ray.rllib.utils.framework import try_import_torch from ray.rllib.utils.typing import TensorType torch, nn = try_import_torch() class TinyAtariCNN(TorchRLModule, ValueFunctionAPI, TargetNetworkAPI): """A tiny CNN stack for fast-learning of Atari envs. The architecture here is the exact same as the one used by the old API stack as CNN default ModelV2. We stack 3 CNN layers based on the config, then a 4th one with linear activation and n 1x1 filters, where n is the number of actions in the (discrete) action space. Simple reshaping (no flattening or extra linear layers necessary) lead to the action logits, which can directly be used inside a distribution or loss. .. testcode:: import numpy as np import gymnasium as gym my_net = TinyAtariCNN( observation_space=gym.spaces.Box(-1.0, 1.0, (42, 42, 4), np.float32), action_space=gym.spaces.Discrete(4), ) B = 10 w = 42 h = 42 c = 4 data = torch.from_numpy( np.random.random_sample(size=(B, w, h, c)).astype(np.float32) ) print(my_net.forward_inference({"obs": data})) print(my_net.forward_exploration({"obs": data})) print(my_net.forward_train({"obs": data})) num_all_params = sum(int(np.prod(p.size())) for p in my_net.parameters()) print(f"num params = {num_all_params}") """ @override(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 flxeible 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..). """ # Get the CNN stack config from our RLModuleConfig's (self.config) # `model_config` property: conv_filters = self.model_config.get("conv_filters") # Default CNN stack with 3 layers: if conv_filters is None: conv_filters = [ [16, 4, 2, "same"], # num filters, kernel wxh, stride wxh, padding type [32, 4, 2, "same"], [256, 11, 1, "valid"], ] # Build the CNN layers. layers = [] # Add user-specified hidden convolutional layers first width, height, in_depth = self.observation_space.shape in_size = [width, height] for filter_specs in conv_filters: if len(filter_specs) == 4: out_depth, kernel_size, strides, padding = filter_specs else: out_depth, kernel_size, strides = filter_specs padding = "same" # Pad like in tensorflow's SAME mode. if padding == "same": padding_size, out_size = same_padding(in_size, kernel_size, strides) layers.append(nn.ZeroPad2d(padding_size)) # No actual padding is performed for "valid" mode, but we will still # compute the output size (input for the next layer). else: out_size = valid_padding(in_size, kernel_size, strides) layer = nn.Conv2d(in_depth, out_depth, kernel_size, strides, bias=True) # Initialize CNN layer kernel and bias. nn.init.xavier_uniform_(layer.weight) nn.init.zeros_(layer.bias) layers.append(layer) # Activation. layers.append(nn.ReLU()) in_size = out_size in_depth = out_depth self._base_cnn_stack = nn.Sequential(*layers) # Add the final CNN 1x1 layer with num_filters == num_actions to be reshaped to # yield the logits (no flattening, no additional linear layers required). _final_conv = nn.Conv2d(in_depth, self.action_space.n, 1, 1, bias=True) nn.init.xavier_uniform_(_final_conv.weight) nn.init.zeros_(_final_conv.bias) self._logits = nn.Sequential( nn.ZeroPad2d(same_padding(in_size, 1, 1)[0]), _final_conv ) self._values = nn.Linear(in_depth, 1) # Mimick old API stack behavior of initializing the value function with `normc` # std=0.01. normc_initializer(0.01)(self._values.weight) @override(TorchRLModule) def _forward(self, batch, **kwargs): # Compute the basic 1D feature tensor (inputs to policy- and value-heads). _, logits = self._compute_embeddings_and_logits(batch) # Return features and logits as ACTION_DIST_INPUTS (categorical distribution). return { Columns.ACTION_DIST_INPUTS: logits, } @override(TorchRLModule) def _forward_train(self, batch, **kwargs): # Compute the basic 1D feature tensor (inputs to policy- and value-heads). embeddings, logits = self._compute_embeddings_and_logits(batch) # Return features and logits as ACTION_DIST_INPUTS (categorical distribution). return { Columns.ACTION_DIST_INPUTS: logits, Columns.EMBEDDINGS: embeddings, } # We implement this RLModule as a TargetNetworkAPI RLModule, so it can be used # by the APPO algorithm. @override(TargetNetworkAPI) def make_target_networks(self) -> None: self._target_base_cnn_stack = make_target_network(self._base_cnn_stack) self._target_logits = make_target_network(self._logits) @override(TargetNetworkAPI) def get_target_network_pairs(self): return [ (self._base_cnn_stack, self._target_base_cnn_stack), (self._logits, self._target_logits), ] @override(TargetNetworkAPI) def forward_target(self, batch, **kw): obs = batch[Columns.OBS].permute(0, 3, 1, 2) embeddings = self._target_base_cnn_stack(obs) logits = self._target_logits(embeddings) return {TARGET_NETWORK_ACTION_DIST_INPUTS: torch.squeeze(logits, dim=[-1, -2])} # We implement this RLModule as a ValueFunctionAPI RLModule, so it can be used # by value-based methods like PPO or IMPALA. @override(ValueFunctionAPI) def compute_values( self, batch: Dict[str, Any], embeddings: Optional[Any] = None, ) -> TensorType: # Features not provided -> We need to compute them first. if embeddings is None: obs = batch[Columns.OBS] embeddings = self._base_cnn_stack(obs.permute(0, 3, 1, 2)) embeddings = torch.squeeze(embeddings, dim=[-1, -2]) return self._values(embeddings).squeeze(-1) def _compute_embeddings_and_logits(self, batch): obs = batch[Columns.OBS].permute(0, 3, 1, 2) embeddings = self._base_cnn_stack(obs) logits = self._logits(embeddings) return ( torch.squeeze(embeddings, dim=[-1, -2]), torch.squeeze(logits, dim=[-1, -2]), )