wehub-resource-sync
130 lines
4.6 KiB
Python
130 lines
4.6 KiB
Python
from typing import Dict
|
|
|
|
import gymnasium as gym
|
|
|
|
from ray.rllib.core import Columns
|
|
from ray.rllib.core.distribution.torch.torch_distribution import (
|
|
TorchCategorical,
|
|
TorchDiagGaussian,
|
|
TorchMultiDistribution,
|
|
)
|
|
from ray.rllib.core.rl_module.apis.value_function_api import ValueFunctionAPI
|
|
from ray.rllib.core.rl_module.rl_module import RLModule
|
|
from ray.rllib.core.rl_module.torch.torch_rl_module import TorchRLModule
|
|
from ray.rllib.utils.annotations import override
|
|
from ray.rllib.utils.framework import try_import_torch
|
|
from ray.rllib.utils.torch_utils import one_hot
|
|
from ray.rllib.utils.typing import TensorType
|
|
|
|
torch, nn = try_import_torch()
|
|
|
|
|
|
class AutoregressiveActionsRLM(TorchRLModule, ValueFunctionAPI):
|
|
"""An RLModule that uses an autoregressive action distribution.
|
|
|
|
Actions are sampled in two steps. The first (prior) action component is sampled from
|
|
a categorical distribution. Then, the second (posterior) action component is sampled
|
|
from a posterior distribution that depends on the first action component and the
|
|
other input data (observations).
|
|
|
|
Note, this RLModule works in combination with any algorithm, whose Learners require
|
|
the `ValueFunctionAPI`.
|
|
"""
|
|
|
|
@override(RLModule)
|
|
def setup(self):
|
|
super().setup()
|
|
|
|
# Assert the action space is correct.
|
|
assert isinstance(self.action_space, gym.spaces.Tuple)
|
|
assert isinstance(self.action_space[0], gym.spaces.Discrete)
|
|
assert self.action_space[0].n == 3
|
|
assert isinstance(self.action_space[1], gym.spaces.Box)
|
|
|
|
self._prior_net = nn.Sequential(
|
|
nn.Linear(
|
|
in_features=self.observation_space.shape[0],
|
|
out_features=256,
|
|
),
|
|
nn.Tanh(),
|
|
nn.Linear(in_features=256, out_features=self.action_space[0].n),
|
|
)
|
|
|
|
self._posterior_net = nn.Sequential(
|
|
nn.Linear(
|
|
in_features=self.observation_space.shape[0] + self.action_space[0].n,
|
|
out_features=256,
|
|
),
|
|
nn.Tanh(),
|
|
nn.Linear(in_features=256, out_features=self.action_space[1].shape[0] * 2),
|
|
)
|
|
|
|
# Build the value function head.
|
|
self._value_net = nn.Sequential(
|
|
nn.Linear(
|
|
in_features=self.observation_space.shape[0],
|
|
out_features=256,
|
|
),
|
|
nn.Tanh(),
|
|
nn.Linear(in_features=256, out_features=1),
|
|
)
|
|
|
|
@override(TorchRLModule)
|
|
def _forward_inference(self, batch: Dict[str, TensorType]) -> Dict[str, TensorType]:
|
|
return self._pi(batch[Columns.OBS], inference=True)
|
|
|
|
@override(TorchRLModule)
|
|
def _forward_exploration(
|
|
self, batch: Dict[str, TensorType], **kwargs
|
|
) -> Dict[str, TensorType]:
|
|
return self._pi(batch[Columns.OBS], inference=False)
|
|
|
|
@override(TorchRLModule)
|
|
def _forward_train(self, batch: Dict[str, TensorType]) -> Dict[str, TensorType]:
|
|
return self._forward_exploration(batch)
|
|
|
|
@override(ValueFunctionAPI)
|
|
def compute_values(self, batch: Dict[str, TensorType], embeddings=None):
|
|
# Value function forward pass.
|
|
vf_out = self._value_net(batch[Columns.OBS])
|
|
# Squeeze out last dimension (single node value head).
|
|
return vf_out.squeeze(-1)
|
|
|
|
# __sphinx_begin__
|
|
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,
|
|
}
|
|
# __sphinx_end__
|
|
|
|
@override(TorchRLModule)
|
|
def get_inference_action_dist_cls(self):
|
|
return TorchMultiDistribution.get_partial_dist_cls(
|
|
child_distribution_cls_struct=(TorchCategorical, TorchDiagGaussian),
|
|
input_lens=(3, 2),
|
|
)
|