[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, {}

rllib/examples/policy/episode_env_aware_policy.py

@@ -0,0 +1,66 @@
+import numpy as np
+
+from ray.rllib.examples.policy.random_policy import RandomPolicy
+from ray.rllib.policy.policy import Policy
+from ray.rllib.policy.sample_batch import SampleBatch
+from ray.rllib.policy.view_requirement import ViewRequirement
+from ray.rllib.utils.annotations import override
+
+
+class EpisodeEnvAwarePolicy(RandomPolicy):
+    """A Policy that always knows the current EpisodeID and EnvID and
+    returns these in its actions."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.episode_id = None
+        self.env_id = None
+
+        class _fake_model:
+            pass
+
+        self.model = _fake_model()
+        self.model.time_major = True
+        self.model.inference_view_requirements = {
+            SampleBatch.EPS_ID: ViewRequirement(),
+            "env_id": ViewRequirement(),
+            SampleBatch.OBS: ViewRequirement(),
+            SampleBatch.PREV_ACTIONS: ViewRequirement(
+                SampleBatch.ACTIONS, space=self.action_space, shift=-1),
+            SampleBatch.PREV_REWARDS: ViewRequirement(
+                SampleBatch.REWARDS, shift=-1),
+        }
+        self.training_view_requirements = dict(
+            **{
+                SampleBatch.NEXT_OBS: ViewRequirement(
+                    SampleBatch.OBS, shift=1),
+                SampleBatch.ACTIONS: ViewRequirement(space=self.action_space),
+                SampleBatch.REWARDS: ViewRequirement(),
+                SampleBatch.DONES: ViewRequirement(),
+            },
+            **self.model.inference_view_requirements)
+
+    @override(Policy)
+    def is_recurrent(self):
+        return True
+
+    @override(Policy)
+    def compute_actions_from_input_dict(self,
+                                        input_dict,
+                                        explore=None,
+                                        timestep=None,
+                                        **kwargs):
+        self.episode_id = input_dict[SampleBatch.EPS_ID][0]
+        self.env_id = input_dict["env_id"][0]
+        # Always return (episodeID, envID)
+        return [
+            np.array([self.episode_id, self.env_id]) for _ in input_dict["obs"]
+        ], [], {}
+
+    @override(Policy)
+    def postprocess_trajectory(self,
+                               sample_batch,
+                               other_agent_batches=None,
+                               episode=None):
+        sample_batch["postprocessed_column"] = sample_batch["obs"] + 1.0
+        return sample_batch