diff --git a/rllib/agents/trainer.py b/rllib/agents/trainer.py index 7d9d3f69d..9055fe378 100644 --- a/rllib/agents/trainer.py +++ b/rllib/agents/trainer.py @@ -848,13 +848,15 @@ class Trainer(Trainable): """ if state is None: state = [] - preprocessed = self.workers.local_worker().preprocessors[ - policy_id].transform(observation) - filtered_obs = self.workers.local_worker().filters[policy_id]( - preprocessed, update=False) + # Check the preprocessor and preprocess, if necessary. + pp = self.workers.local_worker().preprocessors[policy_id] + if type(pp).__name__ != "NoPreprocessor": + observation = pp.transform(observation) + filtered_observation = self.workers.local_worker().filters[policy_id]( + observation, update=False) result = self.get_policy(policy_id).compute_single_action( - filtered_obs, + filtered_observation, state, prev_action, prev_reward, @@ -1091,10 +1093,19 @@ class Trainer(Trainable): @staticmethod def _validate_config(config: PartialTrainerConfigDict): - if not config.get("_use_trajectory_view_api") and \ - config.get("model", {}).get("_time_major"): - raise ValueError("`model._time_major` only supported " - "iff `_use_trajectory_view_api` is True!") + model_config = config.get("model") + if model_config is None: + config["model"] = model_config = {} + + if not config.get("_use_trajectory_view_api"): + traj_view_framestacks = model_config.get("num_framestacks", "auto") + if model_config.get("_time_major"): + raise ValueError("`model._time_major` only supported " + "iff `_use_trajectory_view_api` is True!") + elif traj_view_framestacks != "auto": + raise ValueError("`model.num_framestacks` only supported " + "iff `_use_trajectory_view_api` is True!") + model_config["num_framestacks"] = 0 if isinstance(config["input_evaluation"], tuple): config["input_evaluation"] = list(config["input_evaluation"]) @@ -1104,15 +1115,15 @@ class Trainer(Trainable): config["input_evaluation"])) # Check model config. - prev_a_r = config.get("model", {}).get("lstm_use_prev_action_reward", - DEPRECATED_VALUE) + prev_a_r = model_config.get("lstm_use_prev_action_reward", + DEPRECATED_VALUE) if prev_a_r != DEPRECATED_VALUE: deprecation_warning( "model.lstm_use_prev_action_reward", "model.lstm_use_prev_action and model.lstm_use_prev_reward", error=False) - config["model"]["lstm_use_prev_action"] = prev_a_r - config["model"]["lstm_use_prev_reward"] = prev_a_r + model_config["lstm_use_prev_action"] = prev_a_r + model_config["lstm_use_prev_reward"] = prev_a_r # Check batching/sample collection settings. if config["batch_mode"] not in [ diff --git a/rllib/env/atari_wrappers.py b/rllib/env/atari_wrappers.py index 48ba221b3..1c2b14c77 100644 --- a/rllib/env/atari_wrappers.py +++ b/rllib/env/atari_wrappers.py @@ -226,6 +226,7 @@ class WarpFrame(gym.ObservationWrapper): return frame[:, :, None] +# TODO: (sven) Deprecated class. Remove once traj. view is the norm. class FrameStack(gym.Wrapper): def __init__(self, env, k): """Stack k last frames.""" @@ -255,6 +256,22 @@ class FrameStack(gym.Wrapper): return np.concatenate(self.frames, axis=2) +class FrameStackTrajectoryView(gym.ObservationWrapper): + def __init__(self, env): + """No stacking. Trajectory View API takes care of this.""" + gym.Wrapper.__init__(self, env) + shp = env.observation_space.shape + assert shp[2] == 1 + self.observation_space = spaces.Box( + low=0, + high=255, + shape=(shp[0], shp[1]), + dtype=env.observation_space.dtype) + + def observation(self, observation): + return np.squeeze(observation, axis=-1) + + class ScaledFloatFrame(gym.ObservationWrapper): def __init__(self, env): gym.ObservationWrapper.__init__(self, env) @@ -267,7 +284,12 @@ class ScaledFloatFrame(gym.ObservationWrapper): return np.array(observation).astype(np.float32) / 255.0 -def wrap_deepmind(env, dim=84, framestack=True): +def wrap_deepmind( + env, + dim=84, + # TODO: (sven) Remove once traj. view is norm. + framestack=True, + framestack_via_traj_view_api=False): """Configure environment for DeepMind-style Atari. Note that we assume reward clipping is done outside the wrapper. @@ -286,6 +308,12 @@ def wrap_deepmind(env, dim=84, framestack=True): env = WarpFrame(env, dim) # env = ScaledFloatFrame(env) # TODO: use for dqn? # env = ClipRewardEnv(env) # reward clipping is handled by policy eval - if framestack: + # New way of frame stacking via the trajectory view API (model config key: + # `num_framestacks=[int]`. + if framestack_via_traj_view_api: + env = FrameStackTrajectoryView(env) + # Old way (w/o traj. view API) via model config key: `framestack=True`. + # TODO: (sven) Remove once traj. view is norm. + elif framestack is True: env = FrameStack(env, 4) return env diff --git a/rllib/evaluation/collectors/simple_list_collector.py b/rllib/evaluation/collectors/simple_list_collector.py index 73a6a7619..ce2678193 100644 --- a/rllib/evaluation/collectors/simple_list_collector.py +++ b/rllib/evaluation/collectors/simple_list_collector.py @@ -185,16 +185,17 @@ class _AgentCollector: # each timestep. else: d = np_data[data_col] - # TODO: For now, assume simple 1D data (B x x). - # Will expand this for Atari examples. - assert len(d.shape) == 2 shift_win = view_req.shift_to - view_req.shift_from + 1 data_size = d.itemsize * int(np.product(d.shape[1:])) - + strides = [ + d.itemsize * int(np.product(d.shape[i + 1:])) + for i in range(1, len(d.shape)) + ] data = np.lib.stride_tricks.as_strided( d[self.shift_before - shift_win:], - [self.agent_steps, shift_win, d.shape[1]], - [data_size, data_size, d.itemsize]) + [self.agent_steps, shift_win + ] + [d.shape[i] for i in range(1, len(d.shape))], + [data_size, data_size] + strides) # Set of (probably non-consecutive) indices. # Example: # shift=[-3, 0] @@ -549,6 +550,10 @@ class SimpleListCollector(SampleCollector): input_dict = {} for view_col, view_req in view_reqs.items(): + # Not used for action computations. + if not view_req.used_for_compute_actions: + continue + # Create the batch of data from the different buffers. data_col = view_req.data_col or view_col delta = -1 if data_col in [ diff --git a/rllib/evaluation/rollout_worker.py b/rllib/evaluation/rollout_worker.py index 729105371..4370ee93a 100644 --- a/rllib/evaluation/rollout_worker.py +++ b/rllib/evaluation/rollout_worker.py @@ -5,8 +5,8 @@ import logging import pickle import platform import os -from typing import Callable, Any, List, Dict, Tuple, Union, Optional, \ - TYPE_CHECKING, Type, TypeVar +from typing import Any, Callable, Dict, List, Optional, Tuple, Type, TypeVar, \ + TYPE_CHECKING, Union import ray from ray.rllib.env.atari_wrappers import wrap_deepmind, is_atari @@ -395,11 +395,22 @@ class RolloutWorker(ParallelIteratorWorker): if clip_rewards is None: clip_rewards = True + # framestacking via trajectory view API is enabled. + num_framestacks = model_config.get("num_framestacks", 0) + if not policy_config["_use_trajectory_view_api"]: + model_config["num_framestacks"] = num_framestacks = 0 + elif num_framestacks == "auto": + model_config["num_framestacks"] = num_framestacks = 4 + framestack_traj_view = num_framestacks > 1 + # Deprecated way of framestacking is used. + framestack = model_config.get("framestack") is True + def wrap(env): env = wrap_deepmind( env, dim=model_config.get("dim"), - framestack=model_config.get("framestack")) + framestack=framestack, + framestack_via_traj_view_api=framestack_traj_view) if monitor_path: from gym import wrappers env = wrappers.Monitor(env, monitor_path, resume=True) @@ -1071,27 +1082,33 @@ class RolloutWorker(ParallelIteratorWorker): obs_space = preprocessor.observation_space else: preprocessors[name] = NoPreprocessor(obs_space) - if isinstance(obs_space, gym.spaces.Dict) or \ - isinstance(obs_space, gym.spaces.Tuple): + + if isinstance(obs_space, (gym.spaces.Dict, gym.spaces.Tuple)): raise ValueError( "Found raw Tuple|Dict space as input to policy. " "Please preprocess these observations with a " "Tuple|DictFlatteningPreprocessor.") - if tf1 and tf1.executing_eagerly(): - if hasattr(cls, "as_eager"): - cls = cls.as_eager() - if policy_config.get("eager_tracing"): - cls = cls.with_tracing() - elif not issubclass(cls, TFPolicy): - pass # could be some other type of policy - else: - raise ValueError("This policy does not support eager " - "execution: {}".format(cls)) - if tf1: + # Tf. + framework = policy_config.get("framework", "tf") + if framework in ["tf2", "tf", "tfe"]: + assert tf1 + if framework in ["tf2", "tfe"]: + assert tf1.executing_eagerly() + if hasattr(cls, "as_eager"): + cls = cls.as_eager() + if policy_config.get("eager_tracing"): + cls = cls.with_tracing() + elif not issubclass(cls, TFPolicy): + pass # could be some other type of policy + else: + raise ValueError("This policy does not support eager " + "execution: {}".format(cls)) with tf1.variable_scope(name): policy_map[name] = cls(obs_space, act_space, merged_conf) + # non-tf. else: policy_map[name] = cls(obs_space, act_space, merged_conf) + if self.worker_index == 0: logger.info("Built policy map: {}".format(policy_map)) logger.info("Built preprocessor map: {}".format(preprocessors)) diff --git a/rllib/evaluation/sampler.py b/rllib/evaluation/sampler.py index 0461bcd11..09819cc14 100644 --- a/rllib/evaluation/sampler.py +++ b/rllib/evaluation/sampler.py @@ -17,15 +17,15 @@ from ray.rllib.evaluation.episode import MultiAgentEpisode from ray.rllib.evaluation.rollout_metrics import RolloutMetrics from ray.rllib.evaluation.sample_batch_builder import \ MultiAgentSampleBatchBuilder -from ray.rllib.policy.policy import clip_action, Policy -from ray.rllib.policy.tf_policy import TFPolicy -from ray.rllib.models.preprocessors import Preprocessor -from ray.rllib.utils.filter import Filter from ray.rllib.env.base_env import BaseEnv, ASYNC_RESET_RETURN from ray.rllib.env.atari_wrappers import get_wrapper_by_cls, MonitorEnv +from ray.rllib.models.preprocessors import Preprocessor from ray.rllib.offline import InputReader +from ray.rllib.policy.policy import clip_action, Policy +from ray.rllib.policy.tf_policy import TFPolicy from ray.rllib.utils.annotations import override, DeveloperAPI from ray.rllib.utils.debug import summarize +from ray.rllib.utils.filter import Filter from ray.rllib.utils.numpy import convert_to_numpy from ray.rllib.utils.spaces.space_utils import flatten_to_single_ndarray, \ unbatch @@ -828,13 +828,13 @@ def _process_observations( for agent_id, raw_obs in agent_obs.items(): assert agent_id != "__all__" policy_id: PolicyID = episode.policy_for(agent_id) - prep_obs: EnvObsType = _get_or_raise(preprocessors, - policy_id).transform(raw_obs) + prepr = _get_or_raise(preprocessors, policy_id) + prep_obs: EnvObsType = prepr.transform(raw_obs) if log_once("prep_obs"): logger.info("Preprocessed obs: {}".format(summarize(prep_obs))) - filtered_obs: EnvObsType = _get_or_raise(obs_filters, - policy_id)(prep_obs) + filter = _get_or_raise(obs_filters, policy_id) + filtered_obs: EnvObsType = filter(prep_obs) if log_once("filtered_obs"): logger.info("Filtered obs: {}".format(summarize(filtered_obs))) diff --git a/rllib/models/catalog.py b/rllib/models/catalog.py index 9638ed44b..27be48ea4 100644 --- a/rllib/models/catalog.py +++ b/rllib/models/catalog.py @@ -1,5 +1,6 @@ from functools import partial import gym +from gym.spaces import Box, Dict, Discrete, MultiDiscrete, Tuple import logging import numpy as np import tree @@ -18,7 +19,7 @@ from ray.rllib.models.torch.torch_action_dist import TorchCategorical, \ TorchDeterministic, TorchDiagGaussian, \ TorchMultiActionDistribution, TorchMultiCategorical from ray.rllib.utils.annotations import DeveloperAPI, PublicAPI -from ray.rllib.utils.deprecation import DEPRECATED_VALUE +from ray.rllib.utils.deprecation import DEPRECATED_VALUE, deprecation_warning from ray.rllib.utils.error import UnsupportedSpaceException from ray.rllib.utils.framework import try_import_tf, try_import_torch from ray.rllib.utils.spaces.simplex import Simplex @@ -108,8 +109,17 @@ MODEL_DEFAULTS: ModelConfigDict = { # "attention_use_n_prev_rewards": 0, # == Atari == - # Whether to enable framestack for Atari envs - "framestack": True, + # Which framestacking size to use for Atari envs. + # "auto": Use a value of 4, but only if the env is an Atari env. + # > 1: Use the trajectory view API in the default VisionNets to request the + # last n observations (single, grayscaled 84x84 image frames) as + # inputs. The time axis in the so provided observation tensors + # will come right after the batch axis (channels first format), + # e.g. BxTx84x84, where T=num_framestacks. + # 0 or 1: No framestacking used. + # Use the deprecated `framestack=True`, to disable the above behavor and to + # enable legacy stacking behavior (w/o trajectory view API) instead. + "num_framestacks": "auto", # Final resized frame dimension "dim": 84, # (deprecated) Converts ATARI frame to 1 Channel Grayscale image @@ -134,6 +144,8 @@ MODEL_DEFAULTS: ModelConfigDict = { # Deprecated keys: # Use `lstm_use_prev_action` or `lstm_use_prev_reward` instead. "lstm_use_prev_action_reward": DEPRECATED_VALUE, + # Use `num_framestacks` (int) instead. + "framestack": True, } # __sphinx_doc_end__ # yapf: enable @@ -202,7 +214,7 @@ class ModelCatalog: MultiActionDistribution, TorchMultiActionDistribution): dist_cls = dist_type # Box space -> DiagGaussian OR Deterministic. - elif isinstance(action_space, gym.spaces.Box): + elif isinstance(action_space, Box): if len(action_space.shape) > 1: raise UnsupportedSpaceException( "Action space has multiple dimensions " @@ -218,13 +230,13 @@ class ModelCatalog: dist_cls = TorchDeterministic if framework == "torch" \ else Deterministic # Discrete Space -> Categorical. - elif isinstance(action_space, gym.spaces.Discrete): + elif isinstance(action_space, Discrete): dist_cls = TorchCategorical if framework == "torch" else \ JAXCategorical if framework == "jax" else Categorical # Tuple/Dict Spaces -> MultiAction. elif dist_type in (MultiActionDistribution, TorchMultiActionDistribution) or \ - isinstance(action_space, (gym.spaces.Tuple, gym.spaces.Dict)): + isinstance(action_space, (Tuple, Dict)): return ModelCatalog._get_multi_action_distribution( (MultiActionDistribution if framework == "tf" else TorchMultiActionDistribution), @@ -237,7 +249,7 @@ class ModelCatalog: "Simplex action spaces not supported for torch.") dist_cls = Dirichlet # MultiDiscrete -> MultiCategorical. - elif isinstance(action_space, gym.spaces.MultiDiscrete): + elif isinstance(action_space, MultiDiscrete): dist_cls = TorchMultiCategorical if framework == "torch" else \ MultiCategorical return partial(dist_cls, input_lens=action_space.nvec), \ @@ -265,18 +277,18 @@ class ModelCatalog: """ dl_lib = torch if framework == "torch" else tf - if isinstance(action_space, gym.spaces.Discrete): + if isinstance(action_space, Discrete): return action_space.dtype, (None, ) - elif isinstance(action_space, (gym.spaces.Box, Simplex)): + elif isinstance(action_space, (Box, Simplex)): return dl_lib.float32, (None, ) + action_space.shape - elif isinstance(action_space, gym.spaces.MultiDiscrete): + elif isinstance(action_space, MultiDiscrete): return action_space.dtype, (None, ) + action_space.shape - elif isinstance(action_space, (gym.spaces.Tuple, gym.spaces.Dict)): + elif isinstance(action_space, (Tuple, Dict)): flat_action_space = flatten_space(action_space) size = 0 all_discrete = True for i in range(len(flat_action_space)): - if isinstance(flat_action_space[i], gym.spaces.Discrete): + if isinstance(flat_action_space[i], Discrete): size += 1 else: all_discrete = False @@ -471,7 +483,7 @@ class ModelCatalog: # Try to get a default v2 model. if not model_config.get("custom_model"): v2_class = default_model or ModelCatalog._get_v2_model_class( - obs_space, framework=framework) + obs_space, model_config, framework=framework) if not v2_class: raise ValueError("ModelV2 class could not be determined!") @@ -504,7 +516,7 @@ class ModelCatalog: # Try to get a default v2 model. if not model_config.get("custom_model"): v2_class = default_model or ModelCatalog._get_v2_model_class( - obs_space, framework=framework) + obs_space, model_config, framework=framework) if not v2_class: raise ValueError("ModelV2 class could not be determined!") @@ -536,7 +548,7 @@ class ModelCatalog: elif framework == "jax": v2_class = \ default_model or ModelCatalog._get_v2_model_class( - obs_space, framework=framework) + obs_space, model_config, framework=framework) # Wrap in the requested interface. wrapper = ModelCatalog._wrap_if_needed(v2_class, model_interface) return wrapper(obs_space, action_space, num_outputs, model_config, @@ -661,7 +673,8 @@ class ModelCatalog: @staticmethod def _get_v2_model_class(input_space: gym.Space, - framework: str = "tf") -> ModelV2: + model_config: ModelConfigDict, + framework: str = "tf") -> Type[ModelV2]: VisionNet = None @@ -683,9 +696,13 @@ class ModelCatalog: "framework={} not supported in `ModelCatalog._get_v2_model_" "class`!".format(framework)) - # Discrete/1D obs-spaces. - if isinstance(input_space, gym.spaces.Discrete) or \ - len(input_space.shape) <= 2: + # Discrete/1D obs-spaces or 2D obs space but traj. view framestacking + # disabled. + num_framestacks = model_config.get("num_framestacks", "auto") + if isinstance(input_space, (Discrete, MultiDiscrete)) or \ + len(input_space.shape) == 1 or ( + len(input_space.shape) == 2 and ( + num_framestacks == "auto" or num_framestacks <= 1)): return FCNet # Default Conv2D net. else: @@ -737,3 +754,10 @@ class ModelCatalog: elif config.get("use_lstm"): raise ValueError("`use_lstm` not available for " "framework=jax so far!") + + if config.get("framestack") != DEPRECATED_VALUE: + deprecation_warning( + old="framestack", new="num_framestacks (int)", error=False) + # If old behavior is desired, disable traj. view-style + # framestacking. + config["num_framestacks"] = 0 diff --git a/rllib/models/tf/visionnet.py b/rllib/models/tf/visionnet.py index c2a8de5d2..a6e87a5b5 100644 --- a/rllib/models/tf/visionnet.py +++ b/rllib/models/tf/visionnet.py @@ -4,6 +4,8 @@ import gym from ray.rllib.models.tf.tf_modelv2 import TFModelV2 from ray.rllib.models.tf.misc import normc_initializer from ray.rllib.models.utils import get_activation_fn, get_filter_config +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.policy.view_requirement import ViewRequirement from ray.rllib.utils.framework import try_import_tf from ray.rllib.utils.typing import ModelConfigDict, TensorType @@ -29,9 +31,19 @@ class VisionNetwork(TFModelV2): "Must provide at least 1 entry in `conv_filters`!" no_final_linear = self.model_config.get("no_final_linear") vf_share_layers = self.model_config.get("vf_share_layers") + self.traj_view_framestacking = False - inputs = tf.keras.layers.Input( - shape=obs_space.shape, name="observations") + # Perform Atari framestacking via traj. view API. + if model_config.get("num_framestacks") != "auto" and \ + model_config.get("num_framestacks", 0) > 1: + input_shape = obs_space.shape + (model_config["num_framestacks"], ) + self.data_format = "channels_first" + self.traj_view_framestacking = True + else: + input_shape = obs_space.shape + self.data_format = "channels_last" + + inputs = tf.keras.layers.Input(shape=input_shape, name="observations") last_layer = inputs # Whether the last layer is the output of a Flattened (rather than # a n x (1,1) Conv2D). @@ -142,12 +154,28 @@ class VisionNetwork(TFModelV2): self.base_model = tf.keras.Model(inputs, [conv_out, value_out]) self.register_variables(self.base_model.variables) + # Optional: framestacking obs/new_obs for Atari. + if self.traj_view_framestacking: + from_ = model_config["num_framestacks"] - 1 + self.view_requirements[SampleBatch.OBS].shift = \ + "-{}:0".format(from_) + self.view_requirements[SampleBatch.OBS].shift_from = -from_ + self.view_requirements[SampleBatch.OBS].shift_to = 0 + self.view_requirements[SampleBatch.NEXT_OBS] = ViewRequirement( + data_col=SampleBatch.OBS, + shift="-{}:1".format(from_ - 1), + space=self.view_requirements[SampleBatch.OBS].space, + used_for_compute_actions=False, + ) + def forward(self, input_dict: Dict[str, TensorType], state: List[TensorType], seq_lens: TensorType) -> (TensorType, List[TensorType]): + obs = input_dict["obs"] + if self.data_format == "channels_first": + obs = tf.transpose(obs, [0, 2, 3, 1]) # Explicit cast to float32 needed in eager. - model_out, self._value_out = self.base_model( - tf.cast(input_dict["obs"], tf.float32)) + model_out, self._value_out = self.base_model(tf.cast(obs, tf.float32)) # Our last layer is already flat. if self.last_layer_is_flattened: return model_out, state diff --git a/rllib/models/torch/visionnet.py b/rllib/models/torch/visionnet.py index 1de30d722..cd6352acd 100644 --- a/rllib/models/torch/visionnet.py +++ b/rllib/models/torch/visionnet.py @@ -6,6 +6,8 @@ from ray.rllib.models.torch.torch_modelv2 import TorchModelV2 from ray.rllib.models.torch.misc import normc_initializer, same_padding, \ SlimConv2d, SlimFC from ray.rllib.models.utils import get_filter_config +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.policy.view_requirement import ViewRequirement from ray.rllib.utils.annotations import override from ray.rllib.utils.framework import try_import_torch from ray.rllib.utils.typing import ModelConfigDict, TensorType @@ -19,8 +21,10 @@ class VisionNetwork(TorchModelV2, nn.Module): def __init__(self, obs_space: gym.spaces.Space, action_space: gym.spaces.Space, num_outputs: int, model_config: ModelConfigDict, name: str): + if not model_config.get("conv_filters"): model_config["conv_filters"] = get_filter_config(obs_space.shape) + TorchModelV2.__init__(self, obs_space, action_space, num_outputs, model_config, name) nn.Module.__init__(self) @@ -36,9 +40,18 @@ class VisionNetwork(TorchModelV2, nn.Module): # a n x (1,1) Conv2D). self.last_layer_is_flattened = False self._logits = None + self.traj_view_framestacking = False layers = [] - (w, h, in_channels) = obs_space.shape + # Perform Atari framestacking via traj. view API. + if model_config.get("num_framestacks") != "auto" and \ + model_config.get("num_framestacks", 0) > 1: + (w, h) = obs_space.shape + in_channels = model_config["num_framestacks"] + self.traj_view_framestacking = True + else: + (w, h, in_channels) = obs_space.shape + in_size = [w, h] for out_channels, kernel, stride in filters[:-1]: padding, out_size = same_padding(in_size, kernel, [stride, stride]) @@ -111,7 +124,11 @@ class VisionNetwork(TorchModelV2, nn.Module): activation_fn=None) else: vf_layers = [] - (w, h, in_channels) = obs_space.shape + if self.traj_view_framestacking: + (w, h) = obs_space.shape + in_channels = model_config["num_framestacks"] + else: + (w, h, in_channels) = obs_space.shape in_size = [w, h] for out_channels, kernel, stride in filters[:-1]: padding, out_size = same_padding(in_size, kernel, @@ -150,11 +167,27 @@ class VisionNetwork(TorchModelV2, nn.Module): # Holds the current "base" output (before logits layer). self._features = None + # Optional: framestacking obs/new_obs for Atari. + if self.traj_view_framestacking: + from_ = model_config["num_framestacks"] - 1 + self.view_requirements[SampleBatch.OBS].shift = \ + "-{}:0".format(from_) + self.view_requirements[SampleBatch.OBS].shift_from = -from_ + self.view_requirements[SampleBatch.OBS].shift_to = 0 + self.view_requirements[SampleBatch.NEXT_OBS] = ViewRequirement( + data_col=SampleBatch.OBS, + shift="-{}:1".format(from_ - 1), + space=self.view_requirements[SampleBatch.OBS].space, + ) + @override(TorchModelV2) def forward(self, input_dict: Dict[str, TensorType], state: List[TensorType], seq_lens: TensorType) -> (TensorType, List[TensorType]): - self._features = input_dict["obs"].float().permute(0, 3, 1, 2) + self._features = input_dict["obs"].float() + # No framestacking: + if not self.traj_view_framestacking: + self._features = self._features.permute(0, 3, 1, 2) conv_out = self._convs(self._features) # Store features to save forward pass when getting value_function out. if not self._value_branch_separate: diff --git a/rllib/models/utils.py b/rllib/models/utils.py index ed50ce08c..f866cc944 100644 --- a/rllib/models/utils.py +++ b/rllib/models/utils.py @@ -80,9 +80,11 @@ def get_filter_config(shape): [32, [4, 4], 2], [256, [11, 11], 1], ] - if len(shape) == 3 and shape[:2] == [84, 84]: + if len(shape) in [2, 3] and (shape[:2] == [84, 84] + or shape[1:] == [84, 84]): return filters_84x84 - elif len(shape) == 3 and shape[:2] == [42, 42]: + elif len(shape) in [2, 3] and (shape[:2] == [42, 42] + or shape[1:] == [42, 42]): return filters_42x42 else: raise ValueError( diff --git a/rllib/policy/policy.py b/rllib/policy/policy.py index 8a502e3e4..577ac3d68 100644 --- a/rllib/policy/policy.py +++ b/rllib/policy/policy.py @@ -91,7 +91,9 @@ class Policy(metaclass=ABCMeta): if not hasattr(self, "view_requirements"): self.view_requirements = view_reqs else: - self.view_requirements.update(view_reqs) + for k, v in view_reqs.items(): + if k not in self.view_requirements: + self.view_requirements[k] = v self._model_init_state_automatically_added = False @abstractmethod @@ -546,7 +548,8 @@ class Policy(metaclass=ABCMeta): model=getattr(self, "model", None), num_workers=self.config.get("num_workers", 0), worker_index=self.config.get("worker_index", 0), - framework=getattr(self, "framework", "tf")) + framework=getattr(self, "framework", + self.config.get("framework", "tf"))) return exploration def _get_default_view_requirements(self): diff --git a/rllib/policy/tests/test_trajectory_view_api.py b/rllib/policy/tests/test_trajectory_view_api.py new file mode 100644 index 000000000..e69de29bb diff --git a/rllib/policy/view_requirement.py b/rllib/policy/view_requirement.py index 7d361d8dd..5de9736e5 100644 --- a/rllib/policy/view_requirement.py +++ b/rllib/policy/view_requirement.py @@ -33,6 +33,7 @@ class ViewRequirement: shift: Union[int, str, List[int]] = 0, index: Optional[int] = None, batch_repeat_value: int = 1, + used_for_compute_actions: bool = True, used_for_training: bool = True): """Initializes a ViewRequirement object. @@ -58,6 +59,9 @@ class ViewRequirement: used e.g. for the location of a requested inference dict within the trajectory. Negative values refer to counting from the end of a trajectory. + used_for_compute_actions (bool): Whether the data will be used for + creating input_dicts for `Policy.compute_actions()` calls (or + `Policy.compute_actions_from_input_dict()`). used_for_training (bool): Whether the data will be used for training. If False, the column will not be copied into the final train batch. @@ -81,4 +85,5 @@ class ViewRequirement: self.index = index self.batch_repeat_value = batch_repeat_value + self.used_for_compute_actions = used_for_compute_actions self.used_for_training = used_for_training diff --git a/rllib/utils/test_utils.py b/rllib/utils/test_utils.py index d0263eff5..eda9d1cfa 100644 --- a/rllib/utils/test_utils.py +++ b/rllib/utils/test_utils.py @@ -321,6 +321,12 @@ def check_compute_single_action(trainer, call_kwargs["clip_actions"] = True obs = obs_space.sample() + # Framestacking w/ traj. view API. + framestacks = pol.config["model"].get("num_framestacks", + "auto") + if isinstance(framestacks, int) and framestacks > 1: + obs = np.stack( + [obs] * pol.config["model"]["num_framestacks"]) if isinstance(obs_space, gym.spaces.Box): obs = np.clip(obs, -1.0, 1.0) state_in = None