chore: import upstream snapshot with attribution

rllib/examples/envs/classes/debug_counter_env.py

@@ -0,0 +1,92 @@
+import gymnasium as gym
+import numpy as np
+
+from ray.rllib.env.multi_agent_env import MultiAgentEnv
+
+
+class DebugCounterEnv(gym.Env):
+    """Simple Env that yields a ts counter as observation (0-based).
+
+    Actions have no effect.
+    The episode length is always 15.
+    Reward is always: current ts % 3.
+    """
+
+    def __init__(self, config=None):
+        config = config or {}
+        self.action_space = gym.spaces.Discrete(2)
+        self.observation_space = gym.spaces.Box(0, 100, (1,), dtype=np.float32)
+        self.start_at_t = int(config.get("start_at_t", 0))
+        self.i = self.start_at_t
+
+    def reset(self, *, seed=None, options=None):
+        self.i = self.start_at_t
+        return self._get_obs(), {}
+
+    def step(self, action):
+        self.i += 1
+        terminated = False
+        truncated = self.i >= 15 + self.start_at_t
+        return self._get_obs(), float(self.i % 3), terminated, truncated, {}
+
+    def _get_obs(self):
+        return np.array([self.i], dtype=np.float32)
+
+
+class MultiAgentDebugCounterEnv(MultiAgentEnv):
+    def __init__(self, config):
+        super().__init__()
+        self.num_agents = config["num_agents"]
+        self.base_episode_len = config.get("base_episode_len", 103)
+
+        # Observation dims:
+        # 0=agent ID.
+        # 1=episode ID (0.0 for obs after reset).
+        # 2=env ID (0.0 for obs after reset).
+        # 3=ts (of the agent).
+        self.observation_space = gym.spaces.Dict(
+            {
+                aid: gym.spaces.Box(float("-inf"), float("inf"), (4,))
+                for aid in range(self.num_agents)
+            }
+        )
+
+        # Actions are always:
+        # (episodeID, envID) as floats.
+        self.action_space = gym.spaces.Dict(
+            {
+                aid: gym.spaces.Box(-float("inf"), float("inf"), shape=(2,))
+                for aid in range(self.num_agents)
+            }
+        )
+
+        self.timesteps = [0] * self.num_agents
+        self.terminateds = set()
+        self.truncateds = set()
+
+    def reset(self, *, seed=None, options=None):
+        self.timesteps = [0] * self.num_agents
+        self.terminateds = set()
+        self.truncateds = set()
+        return {
+            i: np.array([i, 0.0, 0.0, 0.0], dtype=np.float32)
+            for i in range(self.num_agents)
+        }, {}
+
+    def step(self, action_dict):
+        obs, rew, terminated, truncated = {}, {}, {}, {}
+        for i, action in action_dict.items():
+            self.timesteps[i] += 1
+            obs[i] = np.array([i, action[0], action[1], self.timesteps[i]])
+            rew[i] = self.timesteps[i] % 3
+            terminated[i] = False
+            truncated[i] = (
+                True if self.timesteps[i] > self.base_episode_len + i else False
+            )
+            if terminated[i]:
+                self.terminateds.add(i)
+            if truncated[i]:
+                self.truncateds.add(i)
+        terminated["__all__"] = len(self.terminateds) == self.num_agents
+        truncated["__all__"] = len(self.truncateds) == self.num_agents
+        return obs, rew, terminated, truncated, {}