wehub-resource-sync
154 lines
6.6 KiB
Python
154 lines
6.6 KiB
Python
from typing import Any, Dict, Optional
|
|
|
|
from ray.rllib.core.columns import Columns
|
|
from ray.rllib.core.distribution.torch.torch_distribution import TorchCategorical
|
|
from ray.rllib.core.rl_module.apis import ValueFunctionAPI
|
|
from ray.rllib.core.rl_module.torch import TorchRLModule
|
|
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()
|
|
|
|
|
|
def _make_categorical_with_temperature(temp):
|
|
"""Helper function to create a new action distribution class.
|
|
|
|
The returned class takes a temperature parameter in its constructor with the default
|
|
value `temp`.
|
|
|
|
Args:
|
|
temp: The default temperature to use for the generated distribution class.
|
|
"""
|
|
|
|
class TorchCategoricalWithTemp(TorchCategorical):
|
|
def __init__(self, logits=None, probs=None, temperature: float = temp):
|
|
"""Initializes a TorchCategoricalWithTemp instance.
|
|
|
|
Args:
|
|
logits: Event log probabilities (non-normalized).
|
|
probs: The probabilities of each event.
|
|
temperature: In case of using logits, this parameter can be used to
|
|
determine the sharpness of the distribution. i.e.
|
|
``probs = softmax(logits / temperature)``. The temperature must be
|
|
strictly positive. A low value (e.g. 1e-10) will result in argmax
|
|
sampling while a larger value will result in uniform sampling.
|
|
"""
|
|
# Either divide logits or probs by the temperature.
|
|
assert (
|
|
temperature > 0.0
|
|
), f"Temperature ({temperature}) must be strictly positive!"
|
|
if logits is not None:
|
|
logits /= temperature
|
|
else:
|
|
probs = torch.nn.functional.softmax(probs / temperature)
|
|
super().__init__(logits, probs)
|
|
|
|
return TorchCategoricalWithTemp
|
|
|
|
|
|
class CustomActionDistributionRLModule(TorchRLModule, ValueFunctionAPI):
|
|
"""A simple TorchRLModule with its own custom action distribution.
|
|
|
|
The distribution differs from the default one by an additional temperature
|
|
parameter applied on top of the Categorical base distribution. See the above
|
|
`TorchCategoricalWithTemp` class for details.
|
|
|
|
.. testcode::
|
|
|
|
import numpy as np
|
|
import gymnasium as gym
|
|
|
|
my_net = CustomActionDistributionRLModule(
|
|
observation_space=gym.spaces.Box(-1.0, 1.0, (4,), np.float32),
|
|
action_space=gym.spaces.Discrete(4),
|
|
model_config={"action_dist_temperature": 5.0},
|
|
)
|
|
|
|
B = 10
|
|
data = torch.from_numpy(
|
|
np.random.random_sample(size=(B, 4)).astype(np.float32)
|
|
)
|
|
# Expect a relatively high-temperature distribution.
|
|
# Set "action_dist_temperature" to small values << 1.0 to approximate greedy
|
|
# behavior (even when stochastically sampling from the distribution).
|
|
print(my_net.forward_exploration({"obs": data}))
|
|
"""
|
|
|
|
@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 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.get("hidden_dim", 256)
|
|
output_dim = self.action_space.n
|
|
|
|
# Define simple encoder, and policy- and vf heads.
|
|
self._encoder = torch.nn.Sequential(
|
|
torch.nn.Linear(input_dim, hidden_dim),
|
|
torch.nn.ReLU(),
|
|
)
|
|
self._policy_net = torch.nn.Linear(hidden_dim, output_dim)
|
|
self._vf = nn.Linear(hidden_dim, 1)
|
|
|
|
# Plug in a custom action dist class.
|
|
# NOTE: If you need more granularity as to which distribution class is used by
|
|
# which forward method (`forward_inference`, `forward_exploration`,
|
|
# `forward_train`), override the RLModule methods
|
|
# `get_inference_action_dist_cls`, `get_exploration_action_dist_cls`, and
|
|
# `get_train_action_dist_cls`, and return
|
|
# your custom class(es) from these. In this case, leave `self.action_dist_cls`
|
|
# set to None, its default value.
|
|
self.action_dist_cls = _make_categorical_with_temperature(
|
|
self.model_config["action_dist_temperature"]
|
|
)
|
|
|
|
@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 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:
|
|
embeddings = self._encoder(batch[Columns.OBS])
|
|
return self._vf(embeddings).squeeze(-1)
|
|
|
|
def _compute_embeddings_and_logits(self, batch):
|
|
embeddings = self._encoder(batch[Columns.OBS])
|
|
logits = self._policy_net(embeddings)
|
|
return embeddings, logits
|