[RLlib] BC/MARWIL/recurrent nets minor cleanups and bug fixes. (#13064)

rllib/agents/dreamer/dreamer.py

@@ -8,8 +8,7 @@ from ray.rllib.agents.dreamer.dreamer_torch_policy import DreamerTorchPolicy
 from ray.rllib.agents.trainer_template import build_trainer
 from ray.rllib.execution.common import STEPS_SAMPLED_COUNTER, \
     LEARNER_INFO, _get_shared_metrics
-from ray.rllib.policy.sample_batch import SampleBatch
-from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID
+from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID, SampleBatch
 from ray.rllib.evaluation.metrics import collect_metrics
 from ray.rllib.agents.dreamer.dreamer_model import DreamerModel
 from ray.rllib.execution.rollout_ops import ParallelRollouts
@@ -215,7 +214,7 @@ class DreamerIteration:
         return frames
 
     def policy_stats(self, fetches):
-        return fetches["default_policy"]["learner_stats"]
+        return fetches[DEFAULT_POLICY_ID]["learner_stats"]
 
 
 def execution_plan(workers, config):

rllib/agents/impala/tests/test_impala.py

@@ -2,6 +2,7 @@ import unittest
 
 import ray
 import ray.rllib.agents.impala as impala
+from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID
 from ray.rllib.utils.framework import try_import_tf
 from ray.rllib.utils.test_utils import check_compute_single_action, \
     framework_iterator
@@ -62,7 +63,7 @@ class TestIMPALA(unittest.TestCase):
         trainer = impala.ImpalaTrainer(config=local_cfg, env="CartPole-v0")
 
         def get_lr(result):
-            return result["info"]["learner"]["default_policy"]["cur_lr"]
+            return result["info"]["learner"][DEFAULT_POLICY_ID]["cur_lr"]
 
         try:
             r1 = trainer.train()

rllib/agents/marwil/bc.py

@@ -10,7 +10,14 @@ from ray.rllib.utils.typing import TrainerConfigDict
 # __sphinx_doc_begin__
 BC_DEFAULT_CONFIG = MARWILTrainer.merge_trainer_configs(
     MARWIL_CONFIG, {
+        # No need to calculate advantages (or do anything else with the
+        # rewards).
         "beta": 0.0,
+        # Advantages (calculated during postprocessing) not important for
+        # behavioral cloning.
+        "postprocess_inputs": False,
+        # No reward estimation.
+        "input_evaluation": [],
     })
 # __sphinx_doc_end__
 # yapf: enable

rllib/agents/marwil/marwil_tf_policy.py

@@ -10,17 +10,35 @@ tf1, tf, tfv = try_import_tf()
 
 
 class ValueNetworkMixin:
-    def __init__(self):
-        @make_tf_callable(self.get_session())
-        def value(ob, prev_action, prev_reward, *state):
-            model_out, _ = self.model({
-                SampleBatch.CUR_OBS: tf.convert_to_tensor([ob]),
-                SampleBatch.PREV_ACTIONS: tf.convert_to_tensor([prev_action]),
-                SampleBatch.PREV_REWARDS: tf.convert_to_tensor([prev_reward]),
-                "is_training": tf.convert_to_tensor(False),
-            }, [tf.convert_to_tensor([s]) for s in state],
-                                      tf.convert_to_tensor([1]))
-            return self.model.value_function()[0]
+    def __init__(self, obs_space, action_space, config):
+
+        # Input dict is provided to us automatically via the Model's
+        # requirements. It's a single-timestep (last one in trajectory)
+        # input_dict.
+        if config["_use_trajectory_view_api"]:
+
+            @make_tf_callable(self.get_session())
+            def value(**input_dict):
+                model_out, _ = self.model.from_batch(
+                    input_dict, is_training=False)
+                # [0] = remove the batch dim.
+                return self.model.value_function()[0]
+
+        # TODO: (sven) Remove once trajectory view API is all-algo default.
+        else:
+
+            @make_tf_callable(self.get_session())
+            def value(ob, prev_action, prev_reward, *state):
+                model_out, _ = self.model({
+                    SampleBatch.CUR_OBS: tf.convert_to_tensor([ob]),
+                    SampleBatch.PREV_ACTIONS: tf.convert_to_tensor(
+                        [prev_action]),
+                    SampleBatch.PREV_REWARDS: tf.convert_to_tensor(
+                        [prev_reward]),
+                    "is_training": tf.convert_to_tensor(False),
+                }, [tf.convert_to_tensor([s]) for s in state],
+                                          tf.convert_to_tensor([1]))
+                return self.model.value_function()[0]
 
         self._value = value
 
@@ -34,51 +52,91 @@ class ValueLoss:
 class ReweightedImitationLoss:
     def __init__(self, policy, state_values, cumulative_rewards, actions,
                  action_dist, beta):
-        # advantage estimation
-        adv = cumulative_rewards - state_values
+        if beta != 0.0:
+            # Advantage Estimation.
+            adv = cumulative_rewards - state_values
 
-        # update averaged advantage norm
-        if policy.config["framework"] in ["tf2", "tfe"]:
-            policy._ma_adv_norm.assign_add(
-                1e-6 *
-                (tf.reduce_mean(tf.math.square(adv)) - policy._ma_adv_norm))
-            # Exponentially weighted advantages.
-            exp_advs = tf.math.exp(beta * tf.math.divide(
-                adv, 1e-8 + tf.math.sqrt(policy._ma_adv_norm)))
-        else:
-            update_adv_norm = tf1.assign_add(
-                ref=policy._ma_adv_norm,
-                value=1e-6 *
-                (tf.reduce_mean(tf.math.square(adv)) - policy._ma_adv_norm))
-
-            # exponentially weighted advantages
-            with tf1.control_dependencies([update_adv_norm]):
+            # Update averaged advantage norm.
+            # Eager.
+            if policy.config["framework"] in ["tf2", "tfe"]:
+                policy._ma_adv_norm.assign_add(1e-6 * (
+                    tf.reduce_mean(tf.math.square(adv)) - policy._ma_adv_norm))
+                # Exponentially weighted advantages.
                 exp_advs = tf.math.exp(beta * tf.math.divide(
                     adv, 1e-8 + tf.math.sqrt(policy._ma_adv_norm)))
+            # Static graph.
+            else:
+                update_adv_norm = tf1.assign_add(
+                    ref=policy._ma_adv_norm,
+                    value=1e-6 * (tf.reduce_mean(tf.math.square(adv)) -
+                                  policy._ma_adv_norm))
 
-        # log\pi_\theta(a|s)
+                # Exponentially weighted advantages.
+                with tf1.control_dependencies([update_adv_norm]):
+                    exp_advs = tf.math.exp(beta * tf.math.divide(
+                        adv, 1e-8 + tf.math.sqrt(policy._ma_adv_norm)))
+            exp_advs = tf.stop_gradient(exp_advs)
+        else:
+            exp_advs = 1.0
+
+        # L = - A * log\pi_\theta(a|s)
         logprobs = action_dist.logp(actions)
-
-        self.loss = -1.0 * tf.reduce_mean(
-            tf.stop_gradient(exp_advs) * logprobs)
+        self.loss = -1.0 * tf.reduce_mean(exp_advs * logprobs)
 
 
 def postprocess_advantages(policy,
                            sample_batch,
                            other_agent_batches=None,
                            episode=None):
-    completed = sample_batch[SampleBatch.DONES][-1]
+    """Postprocesses a trajectory and returns the processed trajectory.
 
-    if completed:
+    The trajectory contains only data from one episode and from one agent.
+    - If  `config.batch_mode=truncate_episodes` (default), sample_batch may
+    contain a truncated (at-the-end) episode, in case the
+    `config.rollout_fragment_length` was reached by the sampler.
+    - If `config.batch_mode=complete_episodes`, sample_batch will contain
+    exactly one episode (no matter how long).
+    New columns can be added to sample_batch and existing ones may be altered.
+
+    Args:
+        policy (Policy): The Policy used to generate the trajectory
+            (`sample_batch`)
+        sample_batch (SampleBatch): The SampleBatch to postprocess.
+        other_agent_batches (Optional[Dict[PolicyID, SampleBatch]]): Optional
+            dict of AgentIDs mapping to other agents' trajectory data (from the
+            same episode). NOTE: The other agents use the same policy.
+        episode (Optional[MultiAgentEpisode]): Optional multi-agent episode
+            object in which the agents operated.
+
+    Returns:
+        SampleBatch: The postprocessed, modified SampleBatch (or a new one).
+    """
+
+    # Trajectory is actually complete -> last r=0.0.
+    if sample_batch[SampleBatch.DONES][-1]:
         last_r = 0.0
+    # Trajectory has been truncated -> last r=VF estimate of last obs.
     else:
-        next_state = []
-        for i in range(policy.num_state_tensors()):
-            next_state.append([sample_batch["state_out_{}".format(i)][-1]])
-        last_r = policy._value(sample_batch[SampleBatch.NEXT_OBS][-1],
-                               sample_batch[SampleBatch.ACTIONS][-1],
-                               sample_batch[SampleBatch.REWARDS][-1],
-                               *next_state)
+        # Input dict is provided to us automatically via the Model's
+        # requirements. It's a single-timestep (last one in trajectory)
+        # input_dict.
+        if policy.config["_use_trajectory_view_api"]:
+            # Create an input dict according to the Model's requirements.
+            index = "last" if SampleBatch.NEXT_OBS in sample_batch.data else -1
+            input_dict = policy.model.get_input_dict(sample_batch, index=index)
+            last_r = policy._value(**input_dict)
+        # TODO: (sven) Remove once trajectory view API is all-algo default.
+        else:
+            next_state = []
+            for i in range(policy.num_state_tensors()):
+                next_state.append(sample_batch["state_out_{}".format(i)][-1])
+            last_r = policy._value(sample_batch[SampleBatch.NEXT_OBS][-1],
+                                   sample_batch[SampleBatch.ACTIONS][-1],
+                                   sample_batch[SampleBatch.REWARDS][-1],
+                                   *next_state)
+
+    # Adds the policy logits, VF preds, and advantages to the batch,
+    # using GAE ("generalized advantage estimation") or not.
     return compute_advantages(
         sample_batch,
         last_r,
@@ -131,7 +189,7 @@ def stats(policy, train_batch):
 
 
 def setup_mixins(policy, obs_space, action_space, config):
-    ValueNetworkMixin.__init__(policy)
+    ValueNetworkMixin.__init__(policy, obs_space, action_space, config)
     # Set up a tf-var for the moving avg (do this here to make it work with
     # eager mode).
     policy._ma_adv_norm = get_variable(

rllib/agents/marwil/marwil_torch_policy.py

@@ -10,18 +10,33 @@ torch, _ = try_import_torch()
 
 
 class ValueNetworkMixin:
-    def __init__(self):
-        def value(ob, prev_action, prev_reward, *state):
-            model_out, _ = self.model({
-                SampleBatch.CUR_OBS: torch.Tensor([ob]).to(self.device),
-                SampleBatch.PREV_ACTIONS: torch.Tensor([prev_action]).to(
-                    self.device),
-                SampleBatch.PREV_REWARDS: torch.Tensor([prev_reward]).to(
-                    self.device),
-                "is_training": False,
-            }, [torch.Tensor([s]).to(self.device) for s in state],
-                                      torch.Tensor([1]).to(self.device))
-            return self.model.value_function()[0]
+    def __init__(self, obs_space, action_space, config):
+
+        # Input dict is provided to us automatically via the Model's
+        # requirements. It's a single-timestep (last one in trajectory)
+        # input_dict.
+        if config["_use_trajectory_view_api"]:
+
+            def value(**input_dict):
+                input_dict = self._lazy_tensor_dict(input_dict)
+                model_out, _ = self.model.from_batch(
+                    input_dict, is_training=False)
+                # [0] = remove the batch dim.
+                return self.model.value_function()[0]
+
+        else:
+
+            def value(ob, prev_action, prev_reward, *state):
+                model_out, _ = self.model({
+                    SampleBatch.CUR_OBS: torch.Tensor([ob]).to(self.device),
+                    SampleBatch.PREV_ACTIONS: torch.Tensor([prev_action]).to(
+                        self.device),
+                    SampleBatch.PREV_REWARDS: torch.Tensor([prev_reward]).to(
+                        self.device),
+                    "is_training": False,
+                }, [torch.Tensor([s]).to(self.device) for s in state],
+                                          torch.Tensor([1]).to(self.device))
+                return self.model.value_function()[0]
 
         self._value = value
 
@@ -72,7 +87,7 @@ def setup_mixins(policy, obs_space, action_space, config):
     policy.ma_adv_norm = torch.tensor(
         [100.0], dtype=torch.float32, requires_grad=False).to(policy.device)
     # Setup Value branch of our NN.
-    ValueNetworkMixin.__init__(policy)
+    ValueNetworkMixin.__init__(policy, obs_space, action_space, config)
 
 
 MARWILTorchPolicy = build_policy_class(

rllib/agents/marwil/tests/test_bc.py

@@ -23,7 +23,8 @@ class TestBC(unittest.TestCase):
     def test_bc_compilation_and_learning_from_offline_file(self):
         """Test whether a BCTrainer can be built with all frameworks.
 
-        And learns from a historic-data file.
+        And learns from a historic-data file (while being evaluated on an
+        actual env using evaluation_num_workers > 0).
         """
         rllib_dir = Path(__file__).parent.parent.parent.parent
         print("rllib dir={}".format(rllib_dir))
@@ -34,7 +35,6 @@ class TestBC(unittest.TestCase):
         config = marwil.BC_DEFAULT_CONFIG.copy()
         config["num_workers"] = 0  # Run locally.
         config["evaluation_num_workers"] = 1
-        config["evaluation_interval"] = 1
         # Evaluate on actual environment.
         config["evaluation_config"] = {"input": "sampler"}
         # Learn from offline data.

rllib/agents/ppo/ppo.py

@@ -21,6 +21,7 @@ from ray.rllib.execution.rollout_ops import ParallelRollouts, ConcatBatches, \
 from ray.rllib.execution.train_ops import TrainOneStep, TrainTFMultiGPU
 from ray.rllib.execution.metric_ops import StandardMetricsReporting
 from ray.rllib.policy.policy import Policy
+from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID
 from ray.rllib.utils.typing import TrainerConfigDict
 from ray.util.iter import LocalIterator
 
@@ -194,10 +195,10 @@ def warn_about_bad_reward_scales(config, result):
 
     # Warn about excessively high VF loss.
     learner_stats = result["info"]["learner"]
-    if "default_policy" in learner_stats:
+    if DEFAULT_POLICY_ID in learner_stats:
         scaled_vf_loss = (config["vf_loss_coeff"] *
-                          learner_stats["default_policy"]["vf_loss"])
-        policy_loss = learner_stats["default_policy"]["policy_loss"]
+                          learner_stats[DEFAULT_POLICY_ID]["vf_loss"])
+        policy_loss = learner_stats[DEFAULT_POLICY_ID]["policy_loss"]
         if config["vf_share_layers"] and scaled_vf_loss > 100:
             logger.warning(
                 "The magnitude of your value function loss is extremely large "

rllib/agents/ppo/tests/test_ddppo.py

@@ -2,6 +2,7 @@ import unittest
 
 import ray
 import ray.rllib.agents.ppo as ppo
+from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID
 from ray.rllib.utils.test_utils import check_compute_single_action, \
     framework_iterator
 
@@ -39,7 +40,7 @@ class TestDDPPO(unittest.TestCase):
             trainer = ppo.ddppo.DDPPOTrainer(config=config, env="CartPole-v0")
             for _ in range(num_iterations):
                 result = trainer.train()
-                lr = result["info"]["learner"]["default_policy"]["cur_lr"]
+                lr = result["info"]["learner"][DEFAULT_POLICY_ID]["cur_lr"]
             trainer.stop()
             assert lr == 0.0, "lr should anneal to 0.0"
 

rllib/agents/trainer.py

@@ -580,7 +580,19 @@ class Trainer(Trainable):
                         pybullet_envs.getList()
                     except (ModuleNotFoundError, ImportError):
                         pass
-                    return gym.make(env, **env_context)
+                    # Try creating a gym env. If this fails we can output a
+                    # decent error message.
+                    try:
+                        return gym.make(env, **env_context)
+                    except gym.error.Error:
+                        raise ValueError(
+                            "The env string you provided ({}) is a) not a "
+                            "known gym/PyBullet environment specifier or b) "
+                            "not registered! To register your custom envs, "
+                            "do `from ray import tune; tune.register('[name]',"
+                            " lambda cfg: [return actual "
+                            "env from here using cfg])`. Then you can use "
+                            "[name] as your config['env'].".format(env))
 
                 self.env_creator = _creator
         else:

rllib/evaluation/rollout_worker.py

@@ -670,7 +670,7 @@ class RolloutWorker(ParallelIteratorWorker):
         # for better compression inside the writer.
         self.output_writer.write(batch)
 
-        # Do off-policy estimation if needed
+        # Do off-policy estimation, if needed.
         if self.reward_estimators:
             for sub_batch in batch.split_by_episode():
                 for estimator in self.reward_estimators:

rllib/evaluation/tests/test_trajectory_view_api.py

@@ -17,7 +17,7 @@ from ray.rllib.examples.policy.episode_env_aware_policy import \
     EpisodeEnvAwareAttentionPolicy, EpisodeEnvAwareLSTMPolicy
 from ray.rllib.models.tf.attention_net import GTrXLNet
 from ray.rllib.policy.rnn_sequencing import pad_batch_to_sequences_of_same_size
-from ray.rllib.policy.sample_batch import SampleBatch
+from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID, SampleBatch
 from ray.rllib.policy.view_requirement import ViewRequirement
 from ray.rllib.utils.annotations import override
 from ray.rllib.utils.test_utils import framework_iterator, check
@@ -288,11 +288,11 @@ class TestTrajectoryViewAPI(unittest.TestCase):
         )
         # Add the next action to the view reqs of the policy.
         # This should be visible then in postprocessing and train batches.
-        rollout_worker_w_api.policy_map["default_policy"].view_requirements[
+        rollout_worker_w_api.policy_map[DEFAULT_POLICY_ID].view_requirements[
             "next_actions"] = ViewRequirement(
                 SampleBatch.ACTIONS, shift=1, space=action_space)
         # Make sure, we have DONEs as well.
-        rollout_worker_w_api.policy_map["default_policy"].view_requirements[
+        rollout_worker_w_api.policy_map[DEFAULT_POLICY_ID].view_requirements[
             "dones"] = ViewRequirement()
         batch = rollout_worker_w_api.sample()
         self.assertTrue("next_actions" in batch.data)

rllib/examples/custom_keras_model.py

@@ -11,6 +11,7 @@ from ray.rllib.models import ModelCatalog
 from ray.rllib.models.tf.misc import normc_initializer
 from ray.rllib.models.tf.tf_modelv2 import TFModelV2
 from ray.rllib.models.tf.visionnet import VisionNetwork as MyVisionNetwork
+from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID
 from ray.rllib.utils.framework import try_import_tf
 
 tf1, tf, tfv = try_import_tf()
@@ -107,8 +108,8 @@ if __name__ == "__main__":
     # Tests https://github.com/ray-project/ray/issues/7293
     def check_has_custom_metric(result):
         r = result["result"]["info"]["learner"]
-        if "default_policy" in r:
-            r = r["default_policy"]
+        if DEFAULT_POLICY_ID in r:
+            r = r[DEFAULT_POLICY_ID]
         assert r["model"]["foo"] == 42, result
 
     if args.run == "DQN":

rllib/examples/rollout_worker_custom_workflow.py

@@ -15,7 +15,7 @@ from ray import tune
 from ray.rllib.evaluation import RolloutWorker
 from ray.rllib.evaluation.metrics import collect_metrics
 from ray.rllib.policy.policy import Policy
-from ray.rllib.policy.sample_batch import SampleBatch
+from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID, SampleBatch
 
 parser = argparse.ArgumentParser()
 parser.add_argument("--gpu", action="store_true")
@@ -76,7 +76,7 @@ def training_workflow(config, reporter):
 
     for _ in range(config["num_iters"]):
         # Broadcast weights to the policy evaluation workers
-        weights = ray.put({"default_policy": policy.get_weights()})
+        weights = ray.put({DEFAULT_POLICY_ID: policy.get_weights()})
         for w in workers:
             w.set_weights.remote(weights)
 

rllib/models/catalog.py

@@ -316,6 +316,14 @@ class ModelCatalog:
                                                      model_interface)
 
             if framework in ["tf2", "tf", "tfe"]:
+                # Try wrapping custom model with LSTM, if required.
+                if model_config.get("use_lstm"):
+                    wrapped_cls = model_cls
+                    forward = wrapped_cls.forward
+                    model_cls = ModelCatalog._wrap_if_needed(
+                        wrapped_cls, LSTMWrapper)
+                    model_cls._wrapped_forward = forward
+
                 # Track and warn if vars were created but not registered.
                 created = set()
 

rllib/models/preprocessors.py

@@ -11,7 +11,10 @@ from ray.rllib.utils.typing import TensorType
 
 ATARI_OBS_SHAPE = (210, 160, 3)
 ATARI_RAM_OBS_SHAPE = (128, )
-VALIDATION_INTERVAL = 100
+
+# Only validate env observations vs the observation space every n times in a
+# Preprocessor.
+OBS_VALIDATION_INTERVAL = 100
 
 logger = logging.getLogger(__name__)
 
@@ -54,7 +57,7 @@ class Preprocessor:
 
     def check_shape(self, observation: Any) -> None:
         """Checks the shape of the given observation."""
-        if self._i % VALIDATION_INTERVAL == 0:
+        if self._i % OBS_VALIDATION_INTERVAL == 0:
             if type(observation) is list and isinstance(
                     self._obs_space, gym.spaces.Box):
                 observation = np.array(observation)
@@ -218,7 +221,7 @@ class TupleFlatteningPreprocessor(Preprocessor):
     @override(Preprocessor)
     def transform(self, observation: TensorType) -> np.ndarray:
         self.check_shape(observation)
-        array = np.zeros(self.shape)
+        array = np.zeros(self.shape, dtype=np.float32)
         self.write(observation, array, 0)
         return array
 
@@ -252,7 +255,7 @@ class DictFlatteningPreprocessor(Preprocessor):
     @override(Preprocessor)
     def transform(self, observation: TensorType) -> np.ndarray:
         self.check_shape(observation)
-        array = np.zeros(self.shape)
+        array = np.zeros(self.shape, dtype=np.float32)
         self.write(observation, array, 0)
         return array
 

rllib/models/tf/recurrent_net.py

@@ -119,6 +119,12 @@ class LSTMWrapper(RecurrentNetwork):
 
         super(LSTMWrapper, self).__init__(obs_space, action_space, None,
                                           model_config, name)
+        # At this point, self.num_outputs is the number of nodes coming
+        # from the wrapped (underlying) model. In other words, self.num_outputs
+        # is the input size for the LSTM layer.
+        # If None, set it to the observation space.
+        if self.num_outputs is None:
+            self.num_outputs = int(np.product(self.obs_space.shape))
 
         self.cell_size = model_config["lstm_cell_size"]
         self.use_prev_action = model_config["lstm_use_prev_action"]
@@ -127,7 +133,7 @@ class LSTMWrapper(RecurrentNetwork):
         if isinstance(action_space, Discrete):
             self.action_dim = action_space.n
         elif isinstance(action_space, MultiDiscrete):
-            self.action_dim = np.product(action_space.nvec)
+            self.action_dim = np.sum(action_space.nvec)
         elif action_space.shape is not None:
             self.action_dim = int(np.product(action_space.shape))
         else:
@@ -143,6 +149,8 @@ class LSTMWrapper(RecurrentNetwork):
         input_layer = tf.keras.layers.Input(
             shape=(None, self.num_outputs), name="inputs")
 
+        # Set self.num_outputs to the number of output nodes desired by the
+        # caller of this constructor.
         self.num_outputs = num_outputs
 
         state_in_h = tf.keras.layers.Input(shape=(self.cell_size, ), name="h")

rllib/models/torch/recurrent_net.py

@@ -118,6 +118,13 @@ class LSTMWrapper(RecurrentNetwork, nn.Module):
         nn.Module.__init__(self)
         super().__init__(obs_space, action_space, None, model_config, name)
 
+        # At this point, self.num_outputs is the number of nodes coming
+        # from the wrapped (underlying) model. In other words, self.num_outputs
+        # is the input size for the LSTM layer.
+        # If None, set it to the observation space.
+        if self.num_outputs is None:
+            self.num_outputs = int(np.product(self.obs_space.shape))
+
         self.cell_size = model_config["lstm_cell_size"]
         self.time_major = model_config.get("_time_major", False)
         self.use_prev_action = model_config["lstm_use_prev_action"]
@@ -126,7 +133,7 @@ class LSTMWrapper(RecurrentNetwork, nn.Module):
         if isinstance(action_space, Discrete):
             self.action_dim = action_space.n
         elif isinstance(action_space, MultiDiscrete):
-            self.action_dim = np.product(action_space.nvec)
+            self.action_dim = np.sum(action_space.nvec)
         elif action_space.shape is not None:
             self.action_dim = int(np.product(action_space.shape))
         else:
@@ -138,9 +145,13 @@ class LSTMWrapper(RecurrentNetwork, nn.Module):
         if self.use_prev_reward:
             self.num_outputs += 1
 
+        # Define actual LSTM layer (with num_outputs being the nodes coming
+        # from the wrapped (underlying) layer).
         self.lstm = nn.LSTM(
             self.num_outputs, self.cell_size, batch_first=not self.time_major)
 
+        # Set self.num_outputs to the number of output nodes desired by the
+        # caller of this constructor.
         self.num_outputs = num_outputs
 
         # Postprocess LSTM output with another hidden layer and compute values.

rllib/offline/io_context.py

@@ -15,7 +15,7 @@ class IOContext:
         config (dict): Configuration of the agent.
         worker_index (int): When there are multiple workers created, this
             uniquely identifies the current worker.
-        worker (RolloutWorker): rollout worker object reference.
+        worker (RolloutWorker): RolloutWorker object reference.
     """
 
     @PublicAPI

rllib/offline/json_reader.py

@@ -12,8 +12,8 @@ except ImportError:
 
 from ray.rllib.offline.input_reader import InputReader
 from ray.rllib.offline.io_context import IOContext
-from ray.rllib.policy.sample_batch import MultiAgentBatch, SampleBatch, \
-    DEFAULT_POLICY_ID
+from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID, MultiAgentBatch, \
+    SampleBatch
 from ray.rllib.utils.annotations import override, PublicAPI
 from ray.rllib.utils.compression import unpack_if_needed
 from ray.rllib.utils.typing import FileType, SampleBatchType
@@ -42,6 +42,10 @@ class JsonReader(InputReader):
         """
 
         self.ioctx = ioctx or IOContext()
+        self.default_policy = None
+        if self.ioctx.worker is not None:
+            self.default_policy = \
+                self.ioctx.worker.policy_map.get(DEFAULT_POLICY_ID)
         if isinstance(inputs, str):
             inputs = os.path.abspath(os.path.expanduser(inputs))
             if os.path.isdir(inputs):
@@ -88,8 +92,8 @@ class JsonReader(InputReader):
         if isinstance(batch, SampleBatch):
             out = []
             for sub_batch in batch.split_by_episode():
-                out.append(self.ioctx.worker.policy_map[DEFAULT_POLICY_ID]
-                           .postprocess_trajectory(sub_batch))
+                out.append(
+                    self.default_policy.postprocess_trajectory(sub_batch))
             return SampleBatch.concat_samples(out)
         else:
             # TODO(ekl) this is trickier since the alignments between agent

rllib/offline/json_writer.py

@@ -120,6 +120,6 @@ def _to_json(batch: SampleBatchType, compress_columns: List[str]) -> str:
         out["policy_batches"] = policy_batches
     else:
         out["type"] = "SampleBatch"
-        for k, v in batch.data.items():
+        for k, v in batch.items():
             out[k] = _to_jsonable(v, compress=k in compress_columns)
     return json.dumps(out)

rllib/policy/sample_batch.py

@@ -2,19 +2,16 @@ import collections
 import numpy as np
 import sys
 import itertools
-from typing import Any, Dict, Iterable, List, Optional, Set, Union
+from typing import Dict, Iterable, List, Optional, Set, Union
 
 from ray.rllib.utils.annotations import PublicAPI, DeveloperAPI
 from ray.rllib.utils.compression import pack, unpack, is_compressed
 from ray.rllib.utils.memory import concat_aligned
-from ray.rllib.utils.typing import TensorType
+from ray.rllib.utils.typing import PolicyID, TensorType
 
 # Default policy id for single agent environments
 DEFAULT_POLICY_ID = "default_policy"
 
-# TODO(ekl) reuse the other id def once we fix imports
-PolicyID = Any
-
 
 @PublicAPI
 class SampleBatch:

rllib/tests/test_execution.py

@@ -16,7 +16,7 @@ from ray.rllib.execution.train_ops import TrainOneStep, ComputeGradients, \
     AverageGradients
 from ray.rllib.execution.replay_buffer import LocalReplayBuffer, \
     ReplayActor
-from ray.rllib.policy.sample_batch import SampleBatch
+from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID, SampleBatch
 from ray.util.iter import LocalIterator, from_range
 from ray.util.iter_metrics import SharedMetrics
 
@@ -173,8 +173,8 @@ def test_train_one_step(ray_start_regular_shared):
     b = a.for_each(TrainOneStep(workers))
     batch, stats = next(b)
     assert isinstance(batch, SampleBatch)
-    assert "default_policy" in stats
-    assert "learner_stats" in stats["default_policy"]
+    assert DEFAULT_POLICY_ID in stats
+    assert "learner_stats" in stats[DEFAULT_POLICY_ID]
     counters = a.shared_metrics.get().counters
     assert counters["num_steps_sampled"] == 100, counters
     assert counters["num_steps_trained"] == 100, counters

rllib/tests/test_model_imports.py

@@ -11,6 +11,7 @@ from ray.rllib.models.catalog import ModelCatalog
 from ray.rllib.models.tf.misc import normc_initializer
 from ray.rllib.models.tf.tf_modelv2 import TFModelV2
 from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
+from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID
 from ray.rllib.utils.framework import try_import_tf, try_import_torch
 from ray.rllib.utils.test_utils import check, framework_iterator
 
@@ -99,27 +100,27 @@ class MyTorchModel(TorchModelV2, nn.Module):
         f = h5py.File(import_file)
         self.layer_1.load_state_dict({
             "weight": torch.Tensor(
-                np.transpose(f["layer1"]["default_policy"]["layer1"][
+                np.transpose(f["layer1"][DEFAULT_POLICY_ID]["layer1"][
                     "kernel:0"].value)),
             "bias": torch.Tensor(
                 np.transpose(
-                    f["layer1"]["default_policy"]["layer1"]["bias:0"].value)),
+                    f["layer1"][DEFAULT_POLICY_ID]["layer1"]["bias:0"].value)),
         })
         self.layer_out.load_state_dict({
             "weight": torch.Tensor(
                 np.transpose(
-                    f["out"]["default_policy"]["out"]["kernel:0"].value)),
+                    f["out"][DEFAULT_POLICY_ID]["out"]["kernel:0"].value)),
             "bias": torch.Tensor(
                 np.transpose(
-                    f["out"]["default_policy"]["out"]["bias:0"].value)),
+                    f["out"][DEFAULT_POLICY_ID]["out"]["bias:0"].value)),
         })
         self.value_branch.load_state_dict({
             "weight": torch.Tensor(
                 np.transpose(
-                    f["value"]["default_policy"]["value"]["kernel:0"].value)),
+                    f["value"][DEFAULT_POLICY_ID]["value"]["kernel:0"].value)),
             "bias": torch.Tensor(
                 np.transpose(
-                    f["value"]["default_policy"]["value"]["bias:0"].value)),
+                    f["value"][DEFAULT_POLICY_ID]["value"]["bias:0"].value)),
         })
 
 
@@ -138,13 +139,13 @@ def model_import_test(algo, config, env):
 
         def current_weight(agent):
             if fw == "tf":
-                return agent.get_weights()["default_policy"][
+                return agent.get_weights()[DEFAULT_POLICY_ID][
                     "default_policy/value/kernel"][0]
             elif fw == "torch":
-                return float(agent.get_weights()["default_policy"][
+                return float(agent.get_weights()[DEFAULT_POLICY_ID][
                     "value_branch.weight"][0][0])
             else:
-                return agent.get_weights()["default_policy"][4][0]
+                return agent.get_weights()[DEFAULT_POLICY_ID][4][0]
 
         # Import weights for our custom model from an h5 file.
         weight_before_import = current_weight(agent)

rllib/tests/test_nested_observation_spaces.py

@@ -230,6 +230,18 @@ def to_list(value):
 class DictSpyModel(TFModelV2):
     capture_index = 0
 
+    def __init__(self, obs_space, action_space, num_outputs, model_config,
+                 name):
+        super().__init__(obs_space, action_space, None, model_config, name)
+        # Will only feed in sensors->pos.
+        input_ = tf.keras.layers.Input(
+            shape=self.obs_space.original_space["sensors"]["position"].shape)
+
+        self.num_outputs = num_outputs or 64
+        out = tf.keras.layers.Dense(self.num_outputs)(input_)
+        self._main_layer = tf.keras.models.Model([input_], [out])
+        self.register_variables(self._main_layer.variables)
+
     def forward(self, input_dict, state, seq_lens):
         def spy(pos, front_cam, task):
             # TF runs this function in an isolated context, so we have to use
@@ -251,14 +263,27 @@ class DictSpyModel(TFModelV2):
             stateful=True)
 
         with tf1.control_dependencies([spy_fn]):
-            output = tf1.layers.dense(input_dict["obs"]["sensors"]["position"],
-                                      self.num_outputs)
+            output = self._main_layer(
+                [input_dict["obs"]["sensors"]["position"]])
+
         return output, []
 
 
 class TupleSpyModel(TFModelV2):
     capture_index = 0
 
+    def __init__(self, obs_space, action_space, num_outputs, model_config,
+                 name):
+        super().__init__(obs_space, action_space, None, model_config, name)
+        # Will only feed in 0th index of observation Tuple space.
+        input_ = tf.keras.layers.Input(
+            shape=self.obs_space.original_space[0].shape)
+
+        self.num_outputs = num_outputs or 64
+        out = tf.keras.layers.Dense(self.num_outputs)(input_)
+        self._main_layer = tf.keras.models.Model([input_], [out])
+        self.register_variables(self._main_layer.variables)
+
     def forward(self, input_dict, state, seq_lens):
         def spy(pos, cam, task):
             # TF runs this function in an isolated context, so we have to use