Source code for ray.rllib.utils.exploration.stochastic_sampling

import functools
import gymnasium as gym
import numpy as np
from typing import Optional, Union

from ray.rllib.models.action_dist import ActionDistribution
from ray.rllib.models.modelv2 import ModelV2
from ray.rllib.utils.annotations import OldAPIStack, override
from ray.rllib.utils.exploration.exploration import Exploration
from ray.rllib.utils.exploration.random import Random
from ray.rllib.utils.framework import (
    get_variable,
    try_import_tf,
    try_import_torch,
    TensorType,
)
from ray.rllib.utils.tf_utils import zero_logps_from_actions

tf1, tf, tfv = try_import_tf()
torch, _ = try_import_torch()


[docs]@OldAPIStack class StochasticSampling(Exploration): """An exploration that simply samples from a distribution. The sampling can be made deterministic by passing explore=False into the call to `get_exploration_action`. Also allows for scheduled parameters for the distributions, such as lowering stddev, temperature, etc.. over time. """
[docs] def __init__( self, action_space: gym.spaces.Space, *, framework: str, model: ModelV2, random_timesteps: int = 0, **kwargs ): """Initializes a StochasticSampling Exploration object. Args: action_space: The gym action space used by the environment. framework: One of None, "tf", "torch". model: The ModelV2 used by the owning Policy. random_timesteps: The number of timesteps for which to act completely randomly. Only after this number of timesteps, actual samples will be drawn to get exploration actions. """ assert framework is not None super().__init__(action_space, model=model, framework=framework, **kwargs) # Create the Random exploration module (used for the first n # timesteps). self.random_timesteps = random_timesteps self.random_exploration = Random( action_space, model=self.model, framework=self.framework, **kwargs ) # The current timestep value (tf-var or python int). self.last_timestep = get_variable( np.array(0, np.int64), framework=self.framework, tf_name="timestep", dtype=np.int64, )
@override(Exploration) def get_exploration_action( self, *, action_distribution: ActionDistribution, timestep: Optional[Union[int, TensorType]] = None, explore: bool = True ): if self.framework == "torch": return self._get_torch_exploration_action( action_distribution, timestep, explore ) else: return self._get_tf_exploration_action_op( action_distribution, timestep, explore ) def _get_tf_exploration_action_op(self, action_dist, timestep, explore): ts = self.last_timestep + 1 stochastic_actions = tf.cond( pred=tf.convert_to_tensor(ts < self.random_timesteps), true_fn=lambda: ( self.random_exploration.get_tf_exploration_action_op( action_dist, explore=True )[0] ), false_fn=lambda: action_dist.sample(), ) deterministic_actions = action_dist.deterministic_sample() action = tf.cond( tf.constant(explore) if isinstance(explore, bool) else explore, true_fn=lambda: stochastic_actions, false_fn=lambda: deterministic_actions, ) logp = tf.cond( tf.math.logical_and( explore, tf.convert_to_tensor(ts >= self.random_timesteps) ), true_fn=lambda: action_dist.sampled_action_logp(), false_fn=functools.partial(zero_logps_from_actions, deterministic_actions), ) # Increment `last_timestep` by 1 (or set to `timestep`). if self.framework == "tf2": self.last_timestep.assign_add(1) return action, logp else: assign_op = ( tf1.assign_add(self.last_timestep, 1) if timestep is None else tf1.assign(self.last_timestep, timestep) ) with tf1.control_dependencies([assign_op]): return action, logp def _get_torch_exploration_action( self, action_dist: ActionDistribution, timestep: Union[TensorType, int], explore: Union[TensorType, bool], ): # Set last timestep or (if not given) increase by one. self.last_timestep = ( timestep if timestep is not None else self.last_timestep + 1 ) # Apply exploration. if explore: # Random exploration phase. if self.last_timestep < self.random_timesteps: action, logp = self.random_exploration.get_torch_exploration_action( action_dist, explore=True ) # Take a sample from our distribution. else: action = action_dist.sample() logp = action_dist.sampled_action_logp() # No exploration -> Return deterministic actions. else: action = action_dist.deterministic_sample() logp = torch.zeros_like(action_dist.sampled_action_logp()) return action, logp