[RLlib] Trajectory View API: Atari framestacking. (#13315)

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 [

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

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 [

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))

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)))
 

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

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

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:

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(

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):

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

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