[RLlib] Issue 11974: Traj view API next-action (shift=+1) not working. (#12407)

* WIP.

* Fix and LINT.

rllib/evaluation/collectors/simple_list_collector.py

@@ -131,12 +131,25 @@ class _AgentCollector:
             # -> skip.
             if data_col not in self.buffers:
                 continue
+            # OBS are already shifted by -1 (the initial obs starts one ts
+            # before all other data columns).
             shift = view_req.shift - \
                 (1 if data_col == SampleBatch.OBS else 0)
             if data_col not in np_data:
                 np_data[data_col] = to_float_np_array(self.buffers[data_col])
+            # Shift is exactly 0: Send trajectory as is.
             if shift == 0:
                 data = np_data[data_col][self.shift_before:]
+            # Shift is positive: We still need to 0-pad at the end here.
+            elif shift > 0:
+                data = to_float_np_array(
+                    self.buffers[data_col][self.shift_before + shift:] + [
+                        np.zeros(
+                            shape=view_req.space.shape,
+                            dtype=view_req.space.dtype) for _ in range(shift)
+                    ])
+            # Shift is negative: Shift into the already existing and 0-padded
+            # "before" area of our buffers.
             else:
                 data = np_data[data_col][self.shift_before + shift:shift]
             if len(data) > 0:

rllib/evaluation/tests/test_trajectory_view_api.py

@@ -1,4 +1,5 @@
 import copy
+import gym
 from gym.spaces import Box, Discrete
 import time
 import unittest
@@ -12,6 +13,7 @@ from ray.rllib.examples.policy.episode_env_aware_policy import \
     EpisodeEnvAwarePolicy
 from ray.rllib.policy.rnn_sequencing import pad_batch_to_sequences_of_same_size
 from ray.rllib.policy.sample_batch import SampleBatch
+from ray.rllib.policy.view_requirement import ViewRequirement
 from ray.rllib.utils.test_utils import framework_iterator, check
 
 
@@ -77,13 +79,14 @@ class TestTrajectoryViewAPI(unittest.TestCase):
         config["model"]["use_lstm"] = True
         config["model"]["lstm_use_prev_action_reward"] = True
 
-        for _ in framework_iterator(config, frameworks="torch"):
+        for _ in framework_iterator(config):
             trainer = ppo.PPOTrainer(config, env="CartPole-v0")
             policy = trainer.get_policy()
             view_req_model = policy.model.inference_view_requirements
             view_req_policy = policy.view_requirements
-            assert len(view_req_model) == 5, view_req_model
-            assert len(view_req_policy) == 18, view_req_policy
+            # 7=obs, prev-a + r, 2x state-in, 2x state-out.
+            assert len(view_req_model) == 7, view_req_model
+            assert len(view_req_policy) == 19, view_req_policy
             for key in [
                     SampleBatch.OBS, SampleBatch.ACTIONS, SampleBatch.REWARDS,
                     SampleBatch.DONES, SampleBatch.NEXT_OBS,
@@ -210,6 +213,38 @@ class TestTrajectoryViewAPI(unittest.TestCase):
                             sampler_perf_wo["mean_action_processing_ms"])
             self.assertLess(duration_w, duration_wo)
 
+    def test_traj_view_next_action(self):
+        action_space = Discrete(2)
+        rollout_worker_w_api = RolloutWorker(
+            env_creator=lambda _: gym.make("CartPole-v0"),
+            policy_config=ppo.DEFAULT_CONFIG,
+            rollout_fragment_length=200,
+            policy_spec=ppo.PPOTorchPolicy,
+            policy_mapping_fn=None,
+            num_envs=1,
+        )
+        # Add the next action to the view reqs of the policy.
+        # This should be visible then in postprocessing and train batches.
+        rollout_worker_w_api.policy_map["default_policy"].view_requirements[
+            "next_actions"] = ViewRequirement(
+                SampleBatch.ACTIONS, shift=1, space=action_space)
+        # Make sure, we have DONEs as well.
+        rollout_worker_w_api.policy_map["default_policy"].view_requirements[
+            "dones"] = ViewRequirement()
+        batch = rollout_worker_w_api.sample()
+        self.assertTrue("next_actions" in batch.data)
+        expected_a_ = None  # expected next action
+        for i in range(len(batch["actions"])):
+            a, d, a_ = batch["actions"][i], batch["dones"][i], \
+                       batch["next_actions"][i]
+            if not d and expected_a_ is not None:
+                check(a, expected_a_)
+            elif d:
+                check(a_, 0)
+                expected_a_ = None
+                continue
+            expected_a_ = a_
+
     def test_traj_view_lstm_functionality(self):
         action_space = Box(-float("inf"), float("inf"), shape=(3, ))
         obs_space = Box(float("-inf"), float("inf"), (4, ))