[RLlib] Fix performance and functionality flaws in attention nets (via Trajectory view API).

rllib/agents/ppo/ppo_tf_policy.py

@@ -193,22 +193,25 @@ def postprocess_ppo_gae(
         last_r = 0.0
     # Trajectory has been truncated -> last r=VF estimate of last obs.
     else:
-        # Input dict is provided to us automatically via the Model's
-        # requirements. It's a single-timestep (last one in trajectory)
-        # input_dict.
-        if policy.config["_use_trajectory_view_api"]:
-            # Create an input dict according to the Model's requirements.
-            input_dict = policy.model.get_input_dict(sample_batch, index=-1)
-            last_r = policy._value(**input_dict)
-        # TODO: (sven) Remove once trajectory view API is all-algo default.
+        state_in_view_req = policy.model.inference_view_requirements.get(
+            "state_in_0")
+        # Attention net.
+        if state_in_view_req and state_in_view_req.shift_from is not None:
+            next_state = []
+            for i in range(policy.num_state_tensors()):
+                view_req = policy.model.inference_view_requirements.get(
+                    "state_in_{}".format(i))
+                next_state.append(sample_batch["state_out_{}".format(i)][
+                                  view_req.shift_from:view_req.shift_to + 1])
+        # Everything else.
         else:
             next_state = []
             for i in range(policy.num_state_tensors()):
                 next_state.append(sample_batch["state_out_{}".format(i)][-1])
-            last_r = policy._value(sample_batch[SampleBatch.NEXT_OBS][-1],
-                                   sample_batch[SampleBatch.ACTIONS][-1],
-                                   sample_batch[SampleBatch.REWARDS][-1],
-                                   *next_state)
+        last_r = policy._value(sample_batch[SampleBatch.NEXT_OBS][-1],
+                               sample_batch[SampleBatch.ACTIONS][-1],
+                               sample_batch[SampleBatch.REWARDS][-1],
+                               *next_state)
 
     # Adds the policy logits, VF preds, and advantages to the batch,
     # using GAE ("generalized advantage estimation") or not.
@@ -303,34 +306,19 @@ def __init__(self, obs_space, action_space, config):
         # observation.
         if config["use_gae"]:
 
-            # Input dict is provided to us automatically via the Model's
-            # requirements. It's a single-timestep (last one in trajectory)
-            # input_dict.
-            if config["_use_trajectory_view_api"]:
-
-                @make_tf_callable(self.get_session())
-                def value(**input_dict):
-                    model_out, _ = self.model.from_batch(
-                        input_dict, is_training=False)
-                    # [0] = remove the batch dim.
-                    return self.model.value_function()[0]
-
-            # TODO: (sven) Remove once trajectory view API is all-algo default.
-            else:
-
-                @make_tf_callable(self.get_session())
-                def value(ob, prev_action, prev_reward, *state):
-                    model_out, _ = self.model({
-                        SampleBatch.CUR_OBS: tf.convert_to_tensor([ob]),
-                        SampleBatch.PREV_ACTIONS: tf.convert_to_tensor(
-                            [prev_action]),
-                        SampleBatch.PREV_REWARDS: tf.convert_to_tensor(
-                            [prev_reward]),
-                        "is_training": tf.convert_to_tensor([False]),
-                    }, [tf.convert_to_tensor([s]) for s in state],
-                                              tf.convert_to_tensor([1]))
-                    # [0] = remove the batch dim.
-                    return self.model.value_function()[0]
+            @make_tf_callable(self.get_session())
+            def value(ob, prev_action, prev_reward, *state):
+                model_out, _ = self.model({
+                    SampleBatch.CUR_OBS: tf.convert_to_tensor([ob]),
+                    SampleBatch.PREV_ACTIONS: tf.convert_to_tensor(
+                        [prev_action]),
+                    SampleBatch.PREV_REWARDS: tf.convert_to_tensor(
+                        [prev_reward]),
+                    "is_training": tf.convert_to_tensor([False]),
+                }, [tf.convert_to_tensor([s]) for s in state],
+                                          tf.convert_to_tensor([1]))
+                # [0] = remove the batch dim.
+                return self.model.value_function()[0]
 
         # When not doing GAE, we do not require the value function's output.
         else:

rllib/agents/ppo/ppo_torch_policy.py

@@ -210,36 +210,22 @@ def __init__(self, obs_space, action_space, config):
         # When doing GAE, we need the value function estimate on the
         # observation.
         if config["use_gae"]:
-            # Input dict is provided to us automatically via the Model's
-            # requirements. It's a single-timestep (last one in trajectory)
-            # input_dict.
-            if config["_use_trajectory_view_api"]:
-
-                def value(**input_dict):
-                    model_out, _ = self.model.from_batch(
-                        convert_to_torch_tensor(input_dict, self.device),
-                        is_training=False)
-                    # [0] = remove the batch dim.
-                    return self.model.value_function()[0]
-
-            # TODO: (sven) Remove once trajectory view API is all-algo default.
-            else:
-
-                def value(ob, prev_action, prev_reward, *state):
-                    model_out, _ = self.model({
-                        SampleBatch.CUR_OBS: convert_to_torch_tensor(
-                            np.asarray([ob]), self.device),
-                        SampleBatch.PREV_ACTIONS: convert_to_torch_tensor(
-                            np.asarray([prev_action]), self.device),
-                        SampleBatch.PREV_REWARDS: convert_to_torch_tensor(
-                            np.asarray([prev_reward]), self.device),
-                        "is_training": False,
-                    }, [
-                        convert_to_torch_tensor(np.asarray([s]), self.device)
-                        for s in state
-                    ], convert_to_torch_tensor(np.asarray([1]), self.device))
-                    # [0] = remove the batch dim.
-                    return self.model.value_function()[0]
+
+            def value(ob, prev_action, prev_reward, *state):
+                model_out, _ = self.model({
+                    SampleBatch.CUR_OBS: convert_to_torch_tensor(
+                        np.asarray([ob]), self.device),
+                    SampleBatch.PREV_ACTIONS: convert_to_torch_tensor(
+                        np.asarray([prev_action]), self.device),
+                    SampleBatch.PREV_REWARDS: convert_to_torch_tensor(
+                        np.asarray([prev_reward]), self.device),
+                    "is_training": False,
+                }, [
+                    convert_to_torch_tensor(np.asarray([s]), self.device)
+                    for s in state
+                ], convert_to_torch_tensor(np.asarray([1]), self.device))
+                # [0] = remove the batch dim.
+                return self.model.value_function()[0]
 
         # When not doing GAE, we do not require the value function's output.
         else: