diff --git a/rllib/evaluation/collectors/simple_list_collector.py b/rllib/evaluation/collectors/simple_list_collector.py index 97623b473..f03d38cec 100644 --- a/rllib/evaluation/collectors/simple_list_collector.py +++ b/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: diff --git a/rllib/evaluation/tests/test_trajectory_view_api.py b/rllib/evaluation/tests/test_trajectory_view_api.py index d92d2b499..ae9fd2248 100644 --- a/rllib/evaluation/tests/test_trajectory_view_api.py +++ b/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, ))