12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455 |
- from gymnasium.spaces import Discrete, MultiDiscrete, Space
- from typing import Union, Optional
- from ray.rllib.models.action_dist import ActionDistribution
- from ray.rllib.models.tf.tf_action_dist import Categorical
- from ray.rllib.models.torch.torch_action_dist import TorchCategorical
- from ray.rllib.utils.annotations import override, PublicAPI
- from ray.rllib.utils.exploration.stochastic_sampling import StochasticSampling
- from ray.rllib.utils.framework import TensorType
- @PublicAPI
- class SoftQ(StochasticSampling):
- """Special case of StochasticSampling w/ Categorical and temperature param.
- Returns a stochastic sample from a Categorical parameterized by the model
- output divided by the temperature. Returns the argmax iff explore=False.
- """
- def __init__(
- self,
- action_space: Space,
- *,
- framework: Optional[str],
- temperature: float = 1.0,
- **kwargs
- ):
- """Initializes a SoftQ Exploration object.
- Args:
- action_space: The gym action space used by the environment.
- temperature: The temperature to divide model outputs by
- before creating the Categorical distribution to sample from.
- framework: One of None, "tf", "torch".
- """
- assert isinstance(action_space, (Discrete, MultiDiscrete))
- super().__init__(action_space, framework=framework, **kwargs)
- self.temperature = temperature
- @override(StochasticSampling)
- def get_exploration_action(
- self,
- action_distribution: ActionDistribution,
- timestep: Union[int, TensorType],
- explore: bool = True,
- ):
- cls = type(action_distribution)
- assert issubclass(cls, (Categorical, TorchCategorical))
- # Re-create the action distribution with the correct temperature
- # applied.
- dist = cls(action_distribution.inputs, self.model, temperature=self.temperature)
- # Delegate to super method.
- return super().get_exploration_action(
- action_distribution=dist, timestep=timestep, explore=explore
- )
|