WIP.

rllib/policy/eager_tf_policy.py

@@ -14,7 +14,6 @@
 from ray.rllib.utils import add_mixins, force_list
 from ray.rllib.utils.annotations import override
 from ray.rllib.utils.framework import try_import_tf
-from ray.rllib.utils.tf_ops import convert_to_non_tf_type
 
 tf1, tf, tfv = try_import_tf()
 logger = logging.getLogger(__name__)
@@ -651,10 +650,6 @@ def _lazy_tensor_dict(self, postprocessed_batch):
             postprocessed_batch.set_get_interceptor(_convert_to_tf)
             return postprocessed_batch
 
-        def _lazy_numpy_dict(self, postprocessed_batch):
-            postprocessed_batch.set_get_interceptor(convert_to_non_tf_type)
-            return postprocessed_batch
-
         @classmethod
         def with_tracing(cls):
             return traced_eager_policy(cls)

rllib/policy/policy.py

@@ -624,6 +624,7 @@ def _initialize_loss_from_dummy_batch(
                 self.view_requirements[key] = \
                     ViewRequirement(space=gym.spaces.Box(
                         -1.0, 1.0, shape=value.shape[1:], dtype=value.dtype))
+        self._dummy_batch.set_get_interceptor(None)
         self.exploration.postprocess_trajectory(self, self._dummy_batch)
         postprocessed_batch = self.postprocess_trajectory(self._dummy_batch)
         seq_lens = None

rllib/policy/sample_batch.py

@@ -168,7 +168,8 @@ def copy(self) -> "SampleBatch":
                 k: np.array(v, copy=True) if isinstance(v, np.ndarray) else v
                 for (k, v) in self.data.items()
             },
-            _seq_lens=self.seq_lens)
+            _seq_lens=self.seq_lens,
+            _dont_check_lens=self.dont_check_lens)
         copy_.set_get_interceptor(self.get_interceptor)
         return copy_
 

rllib/policy/torch_policy.py

@@ -599,15 +599,12 @@ def import_model_from_h5(self, import_file: str) -> None:
         return self.model.import_from_h5(import_file)
 
     def _lazy_tensor_dict(self, postprocessed_batch):
+        if not isinstance(postprocessed_batch, SampleBatch):
+            postprocessed_batch = SampleBatch(postprocessed_batch)
         postprocessed_batch.set_get_interceptor(
             functools.partial(convert_to_torch_tensor, device=self.device))
         return postprocessed_batch
 
-    def _lazy_numpy_dict(self, postprocessed_batch):
-        postprocessed_batch.set_get_interceptor(
-            functools.partial(convert_to_non_torch_type))
-        return postprocessed_batch
-
 
 # TODO: (sven) Unify hyperparam annealing procedures across RLlib (tf/torch)
 #   and for all possible hyperparams, not just lr.