[rllib] Convert torch state arrays to tensors during compute actions (#7162)

* convert to tensor

* normalize fix

rllib/agents/trainer.py

@@ -540,8 +540,16 @@ class Trainer(Trainable):
 
         if self.config["normalize_actions"]:
             inner = self.env_creator
-            self.env_creator = (
-                lambda env_config: NormalizeActionWrapper(inner(env_config)))
+
+            def normalize(env):
+                import gym  # soft dependency
+                if not isinstance(env, gym.Env):
+                    raise ValueError(
+                        "Cannot apply NormalizeActionActionWrapper to env of "
+                        "type {}, which does not subclass gym.Env.", type(env))
+                return NormalizeActionWrapper(env)
+
+            self.env_creator = lambda env_config: normalize(inner(env_config))
 
         Trainer._validate_config(self.config)
         log_level = self.config.get("log_level")

rllib/policy/torch_policy.py

@@ -74,7 +74,9 @@ class TorchPolicy(Policy):
                 input_dict[SampleBatch.PREV_ACTIONS] = prev_action_batch
             if prev_reward_batch:
                 input_dict[SampleBatch.PREV_REWARDS] = prev_reward_batch
-            model_out = self.model(input_dict, state_batches, [1])
+            state_batches = [self._convert_to_tensor(s) for s in state_batches]
+            model_out = self.model(input_dict, state_batches,
+                                   self._convert_to_tensor([1]))
             logits, state = model_out
             action_dist = self.dist_class(logits, self.model)
             # Try our Exploration, if any.
@@ -212,18 +214,17 @@ class TorchPolicy(Policy):
 
     def _lazy_tensor_dict(self, postprocessed_batch):
         train_batch = UsageTrackingDict(postprocessed_batch)
-
-        def convert(arr):
-            if torch.is_tensor(arr):
-                return arr.to(self.device)
-            tensor = torch.from_numpy(np.asarray(arr))
-            if tensor.dtype == torch.double:
-                tensor = tensor.float()
-            return tensor.to(self.device)
-
-        train_batch.set_get_interceptor(convert)
+        train_batch.set_get_interceptor(self._convert_to_tensor)
         return train_batch
 
+    def _convert_to_tensor(self, arr):
+        if torch.is_tensor(arr):
+            return arr.to(self.device)
+        tensor = torch.from_numpy(np.asarray(arr))
+        if tensor.dtype == torch.double:
+            tensor = tensor.float()
+        return tensor.to(self.device)
+
     @override(Policy)
     def export_model(self, export_dir):
         """TODO: implement for torch.