[RLlib] Fix use_lstm flag for ModelV2 (w/o ModelV1 wrapping) and add it for PyTorch. (#8734)

rllib/models/catalog.py

@@ -11,15 +11,18 @@ from ray.rllib.models.preprocessors import get_preprocessor
 from ray.rllib.models.tf.fcnet_v1 import FullyConnectedNetwork
 from ray.rllib.models.tf.lstm_v1 import LSTM
 from ray.rllib.models.tf.modelv1_compat import make_v1_wrapper
+from ray.rllib.models.tf.recurrent_net import LSTMWrapper
 from ray.rllib.models.tf.tf_action_dist import Categorical, \
     Deterministic, DiagGaussian, Dirichlet, \
     MultiActionDistribution, MultiCategorical
 from ray.rllib.models.tf.tf_modelv2 import TFModelV2
 from ray.rllib.models.tf.visionnet_v1 import VisionNetwork
-from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
+from ray.rllib.models.torch.recurrent_net import LSTMWrapper as \
+    TorchLSTMWrapper
 from ray.rllib.models.torch.torch_action_dist import TorchCategorical, \
     TorchDeterministic, TorchDiagGaussian, \
     TorchMultiActionDistribution, TorchMultiCategorical
+from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
 from ray.rllib.utils import try_import_tree
 from ray.rllib.utils.annotations import DeveloperAPI, PublicAPI
 from ray.rllib.utils.deprecation import deprecation_warning, DEPRECATED_VALUE
@@ -57,13 +60,13 @@ MODEL_DEFAULTS = {
     "vf_share_layers": True,
 
     # == LSTM ==
-    # Whether to wrap the model with a LSTM
+    # Whether to wrap the model with an LSTM.
     "use_lstm": False,
-    # Max seq len for training the LSTM, defaults to 20
+    # Max seq len for training the LSTM, defaults to 20.
     "max_seq_len": 20,
-    # Size of the LSTM cell
+    # Size of the LSTM cell.
     "lstm_cell_size": 256,
-    # Whether to feed a_{t-1}, r_{t-1} to LSTM
+    # Whether to feed a_{t-1}, r_{t-1} to LSTM.
     "lstm_use_prev_action_reward": False,
     # When using modelv1 models with a modelv2 algorithm, you may have to
     # define the state shape here (e.g., [256, 256]).
@@ -107,8 +110,9 @@ class ModelCatalog:
         >>> observation = prep.transform(raw_observation)
 
         >>> dist_class, dist_dim = ModelCatalog.get_action_dist(
-                env.action_space, {})
-        >>> model = ModelCatalog.get_model(inputs, dist_dim, options)
+        ...     env.action_space, {})
+        >>> model = ModelCatalog.get_model_v2(
+        ...     obs_space, action_space, num_outputs, options)
         >>> dist = dist_class(model.outputs, model)
         >>> action = dist.sample()
     """
@@ -307,6 +311,7 @@ class ModelCatalog:
             else:
                 model_cls = _global_registry.get(RLLIB_MODEL,
                                                  model_config["custom_model"])
+
             # TODO(sven): Hard-deprecate Model(V1).
             if issubclass(model_cls, ModelV2):
                 logger.info("Wrapping {} as {}".format(model_cls,
@@ -374,10 +379,18 @@ class ModelCatalog:
 
         if framework in ["tf", "tfe"]:
             v2_class = None
-            # try to get a default v2 model
+            # 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, model_config, framework=framework)
+
+            if model_config.get("use_lstm"):
+                wrapped_cls = v2_class
+                forward = wrapped_cls.forward
+                v2_class = ModelCatalog._wrap_if_needed(
+                    wrapped_cls, LSTMWrapper)
+                v2_class._wrapped_forward = forward
+
             # fallback to a default v1 model
             if v2_class is None:
                 if tf.executing_eagerly():
@@ -387,7 +400,7 @@ class ModelCatalog:
                         "observation space: {}, use_lstm={}".format(
                             obs_space, model_config.get("use_lstm")))
                 v2_class = make_v1_wrapper(ModelCatalog.get_model)
-            # wrap in the requested interface
+            # Wrap in the requested interface.
             wrapper = ModelCatalog._wrap_if_needed(v2_class, model_interface)
             return wrapper(obs_space, action_space, num_outputs, model_config,
                            name, **model_kwargs)
@@ -395,6 +408,12 @@ class ModelCatalog:
             v2_class = \
                 default_model or ModelCatalog._get_v2_model_class(
                     obs_space, model_config, framework=framework)
+            if model_config.get("use_lstm"):
+                wrapped_cls = v2_class
+                forward = wrapped_cls.forward
+                v2_class = ModelCatalog._wrap_if_needed(
+                    wrapped_cls, TorchLSTMWrapper)
+                v2_class._wrapped_forward = forward
             # Wrap in the requested interface.
             wrapper = ModelCatalog._wrap_if_needed(v2_class, model_interface)
             return wrapper(obs_space, action_space, num_outputs, model_config,

rllib/models/tf/fcnet.py

@@ -21,7 +21,7 @@ class FullyConnectedNetwork(TFModelV2):
         vf_share_layers = model_config.get("vf_share_layers")
         free_log_std = model_config.get("free_log_std")
 
-        # Maybe generate free-floating bias variables for the second half of
+        # Generate free-floating bias variables for the second half of
         # the outputs.
         if free_log_std:
             assert num_outputs % 2 == 0, (
@@ -34,7 +34,10 @@ class FullyConnectedNetwork(TFModelV2):
         # We are using obs_flat, so take the flattened shape as input.
         inputs = tf.keras.layers.Input(
             shape=(np.product(obs_space.shape), ), name="observations")
-        last_layer = layer_out = inputs
+        # Last hidden layer output (before logits outputs).
+        last_layer = inputs
+        # The action distribution outputs.
+        logits_out = None
         i = 1
 
         # Create layers 0 to second-last.
@@ -49,7 +52,7 @@ class FullyConnectedNetwork(TFModelV2):
         # The last layer is adjusted to be of size num_outputs, but it's a
         # layer with activation.
         if no_final_linear and num_outputs:
-            layer_out = tf.keras.layers.Dense(
+            logits_out = tf.keras.layers.Dense(
                 num_outputs,
                 name="fc_out",
                 activation=activation,
@@ -64,7 +67,7 @@ class FullyConnectedNetwork(TFModelV2):
                     activation=activation,
                     kernel_initializer=normc_initializer(1.0))(last_layer)
             if num_outputs:
-                layer_out = tf.keras.layers.Dense(
+                logits_out = tf.keras.layers.Dense(
                     num_outputs,
                     name="fc_out",
                     activation=None,
@@ -72,38 +75,42 @@ class FullyConnectedNetwork(TFModelV2):
             # Adjust num_outputs to be the number of nodes in the last layer.
             else:
                 self.num_outputs = (
-                    [np.product(obs_space.shape)] + hiddens[-1:-1])[-1]
+                    [np.product(obs_space.shape)] + hiddens[-1:])[-1]
 
         # Concat the log std vars to the end of the state-dependent means.
-        if free_log_std:
+        if free_log_std and logits_out is not None:
 
             def tiled_log_std(x):
                 return tf.tile(
                     tf.expand_dims(self.log_std_var, 0), [tf.shape(x)[0], 1])
 
             log_std_out = tf.keras.layers.Lambda(tiled_log_std)(inputs)
-            layer_out = tf.keras.layers.Concatenate(axis=1)(
-                [layer_out, log_std_out])
+            logits_out = tf.keras.layers.Concatenate(axis=1)(
+                [logits_out, log_std_out])
 
+        last_vf_layer = None
         if not vf_share_layers:
-            # build a parallel set of hidden layers for the value net
-            last_layer = inputs
+            # Build a parallel set of hidden layers for the value net.
+            last_vf_layer = inputs
             i = 1
             for size in hiddens:
-                last_layer = tf.keras.layers.Dense(
+                last_vf_layer = tf.keras.layers.Dense(
                     size,
                     name="fc_value_{}".format(i),
                     activation=activation,
-                    kernel_initializer=normc_initializer(1.0))(last_layer)
+                    kernel_initializer=normc_initializer(1.0))(last_vf_layer)
                 i += 1
 
         value_out = tf.keras.layers.Dense(
             1,
             name="value_out",
             activation=None,
-            kernel_initializer=normc_initializer(0.01))(last_layer)
+            kernel_initializer=normc_initializer(0.01))(
+                last_vf_layer if last_vf_layer is not None else last_layer)
 
-        self.base_model = tf.keras.Model(inputs, [layer_out, value_out])
+        self.base_model = tf.keras.Model(
+            inputs, [(logits_out
+                      if logits_out is not None else last_layer), value_out])
         self.register_variables(self.base_model.variables)
 
     def forward(self, input_dict, state, seq_lens):

rllib/models/tf/recurrent_net.py

@@ -1,3 +1,5 @@
+import numpy as np
+
 from ray.rllib.models.modelv2 import ModelV2
 from ray.rllib.models.tf.tf_modelv2 import TFModelV2
 from ray.rllib.policy.rnn_sequencing import add_time_dimension
@@ -94,3 +96,78 @@ class RecurrentNetwork(TFModelV2):
                 ]
         """
         raise NotImplementedError("You must implement this for a RNN model")
+
+
+class LSTMWrapper(RecurrentNetwork):
+    """An LSTM wrapper serving as an interface for ModelV2s that set use_lstm.
+    """
+
+    def __init__(self, obs_space, action_space, num_outputs, model_config,
+                 name):
+
+        super(LSTMWrapper, self).__init__(obs_space, action_space, None,
+                                          model_config, name)
+
+        self.cell_size = model_config["lstm_cell_size"]
+
+        # Define input layers.
+        input_layer = tf.keras.layers.Input(
+            shape=(None, self.num_outputs), name="inputs")
+
+        self.num_outputs = num_outputs
+
+        state_in_h = tf.keras.layers.Input(shape=(self.cell_size, ), name="h")
+        state_in_c = tf.keras.layers.Input(shape=(self.cell_size, ), name="c")
+        seq_in = tf.keras.layers.Input(shape=(), name="seq_in", dtype=tf.int32)
+
+        # Preprocess observation with a hidden layer and send to LSTM cell
+        lstm_out, state_h, state_c = tf.keras.layers.LSTM(
+            self.cell_size,
+            return_sequences=True,
+            return_state=True,
+            name="lstm")(
+                inputs=input_layer,
+                mask=tf.sequence_mask(seq_in),
+                initial_state=[state_in_h, state_in_c])
+
+        # Postprocess LSTM output with another hidden layer and compute values
+        logits = tf.keras.layers.Dense(
+            self.num_outputs,
+            activation=tf.keras.activations.linear,
+            name="logits")(lstm_out)
+        values = tf.keras.layers.Dense(
+            1, activation=None, name="values")(lstm_out)
+
+        # Create the RNN model
+        self._rnn_model = tf.keras.Model(
+            inputs=[input_layer, seq_in, state_in_h, state_in_c],
+            outputs=[logits, values, state_h, state_c])
+        self.register_variables(self._rnn_model.variables)
+        self._rnn_model.summary()
+
+    @override(RecurrentNetwork)
+    def forward(self, input_dict, state, seq_lens):
+        assert seq_lens is not None
+        # Push obs through "unwrapped" net's `forward()` first.
+        wrapped_out, _ = self._wrapped_forward(input_dict, [], None)
+
+        # Then through our LSTM.
+        input_dict["obs_flat"] = wrapped_out
+        return super().forward(input_dict, state, seq_lens)
+
+    @override(RecurrentNetwork)
+    def forward_rnn(self, inputs, state, seq_lens):
+        model_out, self._value_out, h, c = self._rnn_model([inputs, seq_lens] +
+                                                           state)
+        return model_out, [h, c]
+
+    @override(ModelV2)
+    def get_initial_state(self):
+        return [
+            np.zeros(self.cell_size, np.float32),
+            np.zeros(self.cell_size, np.float32),
+        ]
+
+    @override(ModelV2)
+    def value_function(self):
+        return tf.reshape(self._value_out, [-1])

rllib/models/torch/fcnet.py

@@ -13,36 +13,32 @@ torch, nn = try_import_torch()
 logger = logging.getLogger(__name__)
 
 
-class FullyConnectedNetwork(TorchModelV2, nn.Module):
+class FullyConnectedNetwork(TorchModelV2):
     """Generic fully connected network."""
 
     def __init__(self, obs_space, action_space, num_outputs, model_config,
                  name):
         TorchModelV2.__init__(self, obs_space, action_space, num_outputs,
                               model_config, name)
-        nn.Module.__init__(self)
 
         activation = get_activation_fn(
             model_config.get("fcnet_activation"), framework="torch")
         hiddens = model_config.get("fcnet_hiddens")
         no_final_linear = model_config.get("no_final_linear")
+        self.vf_share_layers = model_config.get("vf_share_layers")
         self.free_log_std = model_config.get("free_log_std")
 
-        # TODO(sven): implement case: vf_shared_layers = False.
-        # vf_share_layers = model_config.get("vf_share_layers")
-
-        logger.debug("Constructing fcnet {} {}".format(hiddens, activation))
-        layers = []
-        prev_layer_size = int(np.product(obs_space.shape))
-        self._logits = None
-
-        # Maybe generate free-floating bias variables for the second half of
+        # Generate free-floating bias variables for the second half of
         # the outputs.
         if self.free_log_std:
             assert num_outputs % 2 == 0, (
                 "num_outputs must be divisible by two", num_outputs)
             num_outputs = num_outputs // 2
 
+        layers = []
+        prev_layer_size = int(np.product(obs_space.shape))
+        self._logits = None
+
         # Create layers 0 to second-last.
         for size in hiddens[:-1]:
             layers.append(
@@ -82,15 +78,30 @@ class FullyConnectedNetwork(TorchModelV2, nn.Module):
                     activation_fn=None)
             else:
                 self.num_outputs = (
-                    [np.product(obs_space.shape)] + hiddens[-1:-1])[-1]
+                    [np.product(obs_space.shape)] + hiddens[-1:])[-1]
 
         # Layer to add the log std vars to the state-dependent means.
-        if self.free_log_std:
+        if self.free_log_std and self._logits:
             self._append_free_log_std = AppendBiasLayer(num_outputs)
 
         self._hidden_layers = nn.Sequential(*layers)
 
-        # TODO(sven): Implement non-shared value branch.
+        self._value_branch_separate = None
+        if not self.vf_share_layers:
+            # Build a parallel set of hidden layers for the value net.
+            prev_vf_layer_size = int(np.product(obs_space.shape))
+            self._value_branch_separate = []
+            for size in hiddens:
+                self._value_branch_separate.append(
+                    SlimFC(
+                        in_size=prev_vf_layer_size,
+                        out_size=size,
+                        activation_fn=activation,
+                        initializer=normc_initializer(1.0)))
+                prev_vf_layer_size = size
+            self._value_branch_separate = nn.Sequential(
+                *self._value_branch_separate)
+
         self._value_branch = SlimFC(
             in_size=prev_layer_size,
             out_size=1,
@@ -98,11 +109,14 @@ class FullyConnectedNetwork(TorchModelV2, nn.Module):
             activation_fn=None)
         # Holds the current "base" output (before logits layer).
         self._features = None
+        # Holds the last input, in case value branch is separate.
+        self._last_flat_in = None
 
     @override(TorchModelV2)
     def forward(self, input_dict, state, seq_lens):
         obs = input_dict["obs_flat"].float()
-        self._features = self._hidden_layers(obs.reshape(obs.shape[0], -1))
+        self._last_flat_in = obs.reshape(obs.shape[0], -1)
+        self._features = self._hidden_layers(self._last_flat_in)
         logits = self._logits(self._features) if self._logits else \
             self._features
         if self.free_log_std:
@@ -112,4 +126,8 @@ class FullyConnectedNetwork(TorchModelV2, nn.Module):
     @override(TorchModelV2)
     def value_function(self):
         assert self._features is not None, "must call forward() first"
-        return self._value_branch(self._features).squeeze(1)
+        if self._value_branch_separate:
+            return self._value_branch(
+                self._value_branch_separate(self._last_flat_in)).squeeze(1)
+        else:
+            return self._value_branch(self._features).squeeze(1)

rllib/models/torch/recurrent_net.py

@@ -1,6 +1,7 @@
 import numpy as np
 
 from ray.rllib.models.modelv2 import ModelV2
+from ray.rllib.models.torch.misc import SlimFC
 from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
 from ray.rllib.policy.rnn_sequencing import add_time_dimension
 from ray.rllib.utils.annotations import override, DeveloperAPI
@@ -10,7 +11,7 @@ torch, nn = try_import_torch()
 
 
 @DeveloperAPI
-class RecurrentNetwork(TorchModelV2, nn.Module):
+class RecurrentNetwork(TorchModelV2):
     """Helper class to simplify implementing RNN models with TorchModelV2.
 
     Instead of implementing forward(), you can implement forward_rnn() which
@@ -51,12 +52,6 @@ class RecurrentNetwork(TorchModelV2, nn.Module):
             return q, [h]
     """
 
-    def __init__(self, obs_space, action_space, num_outputs, model_config,
-                 name):
-        TorchModelV2.__init__(self, obs_space, action_space, num_outputs,
-                              model_config, name)
-        nn.Module.__init__(self)
-
     @override(ModelV2)
     def forward(self, input_dict, state, seq_lens):
         """Adds time dimension to batch before sending inputs to forward_rnn().
@@ -90,3 +85,65 @@ class RecurrentNetwork(TorchModelV2, nn.Module):
                 return model_out, [h, c]
         """
         raise NotImplementedError("You must implement this for an RNN model")
+
+
+class LSTMWrapper(RecurrentNetwork):
+    """An LSTM wrapper serving as an interface for ModelV2s that set use_lstm.
+    """
+
+    def __init__(self, obs_space, action_space, num_outputs, model_config,
+                 name):
+
+        super(LSTMWrapper, self).__init__(obs_space, action_space, None,
+                                          model_config, name)
+
+        self.cell_size = model_config["lstm_cell_size"]
+        self.lstm = nn.LSTM(self.num_outputs, self.cell_size, batch_first=True)
+
+        self.num_outputs = num_outputs
+
+        # Postprocess LSTM output with another hidden layer and compute values.
+        self._logits_branch = SlimFC(
+            in_size=self.cell_size,
+            out_size=self.num_outputs,
+            activation_fn=None,
+            initializer=torch.nn.init.xavier_uniform_)
+        self._value_branch = SlimFC(
+            in_size=self.cell_size,
+            out_size=1,
+            activation_fn=None,
+            initializer=torch.nn.init.xavier_uniform_)
+
+    @override(RecurrentNetwork)
+    def forward(self, input_dict, state, seq_lens):
+        assert seq_lens is not None
+        # Push obs through "unwrapped" net's `forward()` first.
+        wrapped_out, _ = self._wrapped_forward(input_dict, [], None)
+
+        # Then through our LSTM.
+        input_dict["obs_flat"] = wrapped_out
+        return super().forward(input_dict, state, seq_lens)
+
+    @override(RecurrentNetwork)
+    def forward_rnn(self, inputs, state, seq_lens):
+        self._features, [h, c] = self.lstm(
+            inputs,
+            [torch.unsqueeze(state[0], 0),
+             torch.unsqueeze(state[1], 0)])
+        model_out = self._logits_branch(self._features)
+        return model_out, [torch.squeeze(h, 0), torch.squeeze(c, 0)]
+
+    @override(ModelV2)
+    def get_initial_state(self):
+        # Place hidden states on same device as model.
+        linear = next(self._logits_branch._model.children())
+        h = [
+            linear.weight.new(1, self.cell_size).zero_().squeeze(0),
+            linear.weight.new(1, self.cell_size).zero_().squeeze(0)
+        ]
+        return h
+
+    @override(ModelV2)
+    def value_function(self):
+        assert self._features is not None, "must call forward() first"
+        return torch.reshape(self._value_branch(self._features), [-1])

rllib/models/torch/torch_modelv2.py

@@ -6,7 +6,7 @@ _, nn = try_import_torch()
 
 
 @PublicAPI
-class TorchModelV2(ModelV2):
+class TorchModelV2(ModelV2, nn.Module):
     """Torch version of ModelV2.
 
     Note that this class by itself is not a valid model unless you
@@ -27,11 +27,6 @@ class TorchModelV2(ModelV2):
                 self._value_branch = ...
         """
 
-        if not isinstance(self, nn.Module):
-            raise ValueError(
-                "Subclasses of TorchModelV2 must also inherit from "
-                "nn.Module, e.g., MyModel(TorchModel, nn.Module)")
-
         ModelV2.__init__(
             self,
             obs_space,
@@ -40,6 +35,7 @@ class TorchModelV2(ModelV2):
             model_config,
             name,
             framework="torch")
+        nn.Module.__init__(self)
 
     @override(ModelV2)
     def variables(self, as_dict=False):

rllib/models/torch/visionnet.py

@@ -9,14 +9,13 @@ from ray.rllib.utils import try_import_torch
 _, nn = try_import_torch()
 
 
-class VisionNetwork(TorchModelV2, nn.Module):
+class VisionNetwork(TorchModelV2):
     """Generic vision network."""
 
     def __init__(self, obs_space, action_space, num_outputs, model_config,
                  name):
         TorchModelV2.__init__(self, obs_space, action_space, num_outputs,
                               model_config, name)
-        nn.Module.__init__(self)
 
         activation = get_activation_fn(
             model_config.get("conv_activation"), framework="torch")