[RLlib] Trajectory view API - 03 Fast LSTM + prev actions/rewards (#9950)

rllib/examples/env/debug_counter_env.py

@@ -1,4 +1,7 @@
 import gym
+import numpy as np
+
+from ray.rllib.env.multi_agent_env import MultiAgentEnv
 
 
 class DebugCounterEnv(gym.Env):
@@ -21,3 +24,43 @@ class DebugCounterEnv(gym.Env):
     def step(self, action):
         self.i += 1
         return [self.i], self.i % 3, self.i >= 15, {}
+
+
+class MultiAgentDebugCounterEnv(MultiAgentEnv):
+    def __init__(self, config):
+        self.num_agents = config["num_agents"]
+        self.p_done = config.get("p_done", 0.02)
+        # Actions are always:
+        # (episodeID, envID) as floats.
+        self.action_space = \
+            gym.spaces.Box(-float("inf"), float("inf"), shape=(2, ))
+        # 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.Box(float("-inf"), float("inf"), (4, ))
+        self.timesteps = [0] * self.num_agents
+        self.dones = set()
+
+    def reset(self):
+        self.dones = 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, done = {}, {}, {}
+        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
+            done[i] = bool(
+                np.random.choice(
+                    [True, False], p=[self.p_done, 1.0 - self.p_done]))
+            if done[i]:
+                self.dones.add(i)
+        done["__all__"] = len(self.dones) == self.num_agents
+        return obs, rew, done, {}