diff --git a/rllib/agents/trainer.py b/rllib/agents/trainer.py index 619caca14..36e308ab1 100644 --- a/rllib/agents/trainer.py +++ b/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") diff --git a/rllib/policy/torch_policy.py b/rllib/policy/torch_policy.py index 4d63af818..819fc3f20 100644 --- a/rllib/policy/torch_policy.py +++ b/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.