chore: import upstream snapshot with attribution

rllib/examples/_old_api_stack/policy/cliff_walking_wall_policy.py

@@ -0,0 +1,101 @@
+# @OldAPIStack
+from typing import Dict, List, Optional, Tuple, Union
+
+import gymnasium as gym
+import numpy as np
+
+from ray.rllib.models.torch.torch_action_dist import TorchCategorical
+from ray.rllib.policy.policy import Policy, ViewRequirement
+from ray.rllib.policy.sample_batch import SampleBatch
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.debug import update_global_seed_if_necessary
+from ray.rllib.utils.typing import AlgorithmConfigDict, TensorStructType, TensorType
+
+
+class CliffWalkingWallPolicy(Policy):
+    """Optimal RLlib policy for the CliffWalkingWallEnv environment, defined in
+    ray/rllib/examples/env/cliff_walking_wall_env.py, with epsilon-greedy exploration.
+
+    The policy takes a random action with probability epsilon, specified
+    by `config["epsilon"]`, and the optimal action with probability  1 - epsilon.
+    """
+
+    @override(Policy)
+    def __init__(
+        self,
+        observation_space: gym.Space,
+        action_space: gym.Space,
+        config: AlgorithmConfigDict,
+    ):
+        update_global_seed_if_necessary(seed=config.get("seed"))
+        super().__init__(observation_space, action_space, config)
+
+        # Known optimal action dist for each of the 48 states and 4 actions
+        self.action_dist = np.zeros((48, 4), dtype=float)
+        # Starting state: go up
+        self.action_dist[36] = (1, 0, 0, 0)
+        # Cliff + Goal: never actually used, set to random
+        self.action_dist[37:] = (0.25, 0.25, 0.25, 0.25)
+        # Row 2; always go right
+        self.action_dist[24:36] = (0, 1, 0, 0)
+        # Row 0 and Row 1; go down or go right
+        self.action_dist[0:24] = (0, 0.5, 0.5, 0)
+        # Col 11; always go down, supercedes previous values
+        self.action_dist[[11, 23, 35]] = (0, 0, 1, 0)
+        assert np.allclose(self.action_dist.sum(-1), 1)
+
+        # Epsilon-Greedy action selection
+        epsilon = config.get("epsilon", 0.0)
+        self.action_dist = self.action_dist * (1 - epsilon) + epsilon / 4
+        assert np.allclose(self.action_dist.sum(-1), 1)
+
+        # Attributes required for RLlib; note that while CliffWalkingWallPolicy
+        # inherits from Policy, it actually implements TorchPolicyV2.
+        self.view_requirements[SampleBatch.ACTION_PROB] = ViewRequirement()
+        self.device = "cpu"
+        self.model = None
+        self.dist_class = TorchCategorical
+
+    @override(Policy)
+    def compute_actions(
+        self,
+        obs_batch: Union[List[TensorStructType], TensorStructType],
+        state_batches: Optional[List[TensorType]] = None,
+        **kwargs,
+    ) -> Tuple[TensorType, List[TensorType], Dict[str, TensorType]]:
+        obs = np.array(obs_batch, dtype=int)
+        action_probs = self.action_dist[obs]
+        actions = np.zeros(len(obs), dtype=int)
+        for i in range(len(obs)):
+            actions[i] = np.random.choice(4, p=action_probs[i])
+        return (
+            actions,
+            [],
+            {SampleBatch.ACTION_PROB: action_probs[np.arange(len(obs)), actions]},
+        )
+
+    @override(Policy)
+    def compute_log_likelihoods(
+        self,
+        actions: Union[List[TensorType], TensorType],
+        obs_batch: Union[List[TensorType], TensorType],
+        **kwargs,
+    ) -> TensorType:
+        obs = np.array(obs_batch, dtype=int)
+        actions = np.array(actions, dtype=int)
+        # Compute action probs for all possible actions
+        action_probs = self.action_dist[obs]
+        # Take the action_probs corresponding to the specified actions
+        action_probs = action_probs[np.arange(len(obs)), actions]
+        # Ignore RuntimeWarning thrown by np.log(0) if action_probs is 0
+        with np.errstate(divide="ignore"):
+            return np.log(action_probs)
+
+    def action_distribution_fn(
+        self, model, obs_batch: TensorStructType, **kwargs
+    ) -> Tuple[TensorType, type, List[TensorType]]:
+        obs = np.array(obs_batch[SampleBatch.OBS], dtype=int)
+        action_probs = self.action_dist[obs]
+        # Ignore RuntimeWarning thrown by np.log(0) if action_probs is 0
+        with np.errstate(divide="ignore"):
+            return np.log(action_probs), TorchCategorical, None

rllib/examples/_old_api_stack/policy/random_policy.py

@@ -0,0 +1,102 @@
+# @OldAPIStack
+import random
+from typing import (
+    List,
+    Optional,
+    Union,
+)
+
+import numpy as np
+import tree  # pip install dm_tree
+from gymnasium.spaces import Box
+
+from ray.rllib.policy.policy import Policy
+from ray.rllib.policy.sample_batch import SampleBatch
+from ray.rllib.utils.annotations import override
+from ray.rllib.utils.typing import ModelWeights, TensorStructType, TensorType
+
+
+class RandomPolicy(Policy):
+    """Hand-coded policy that returns random actions."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        # Whether for compute_actions, the bounds given in action_space
+        # should be ignored (default: False). This is to test action-clipping
+        # and any Env's reaction to bounds breaches.
+        if self.config.get("ignore_action_bounds", False) and isinstance(
+            self.action_space, Box
+        ):
+            self.action_space_for_sampling = Box(
+                -float("inf"),
+                float("inf"),
+                shape=self.action_space.shape,
+                dtype=self.action_space.dtype,
+            )
+        else:
+            self.action_space_for_sampling = self.action_space
+
+    @override(Policy)
+    def init_view_requirements(self):
+        super().init_view_requirements()
+        # Disable for_training and action attributes for SampleBatch.INFOS column
+        # since it can not be properly batched.
+        vr = self.view_requirements[SampleBatch.INFOS]
+        vr.used_for_training = False
+        vr.used_for_compute_actions = False
+
+    @override(Policy)
+    def compute_actions(
+        self,
+        obs_batch: Union[List[TensorStructType], TensorStructType],
+        state_batches: Optional[List[TensorType]] = None,
+        prev_action_batch: Union[List[TensorStructType], TensorStructType] = None,
+        prev_reward_batch: Union[List[TensorStructType], TensorStructType] = None,
+        **kwargs,
+    ):
+        # Alternatively, a numpy array would work here as well.
+        # e.g.: np.array([random.choice([0, 1])] * len(obs_batch))
+        obs_batch_size = len(tree.flatten(obs_batch)[0])
+        return (
+            [self.action_space_for_sampling.sample() for _ in range(obs_batch_size)],
+            [],
+            {},
+        )
+
+    @override(Policy)
+    def learn_on_batch(self, samples):
+        """No learning."""
+        return {}
+
+    @override(Policy)
+    def compute_log_likelihoods(
+        self,
+        actions,
+        obs_batch,
+        state_batches=None,
+        prev_action_batch=None,
+        prev_reward_batch=None,
+        **kwargs,
+    ):
+        return np.array([random.random()] * len(obs_batch))
+
+    @override(Policy)
+    def get_weights(self) -> ModelWeights:
+        """No weights to save."""
+        return {}
+
+    @override(Policy)
+    def set_weights(self, weights: ModelWeights) -> None:
+        """No weights to set."""
+        pass
+
+    @override(Policy)
+    def _get_dummy_batch_from_view_requirements(self, batch_size: int = 1):
+        return SampleBatch(
+            {
+                SampleBatch.OBS: tree.map_structure(
+                    lambda s: s[None], self.observation_space.sample()
+                ),
+            }
+        )