[RLlib] Fix broken example: tf-eager with custom-RNN (#6732). (#7021)

* WIP.

* Fix float32 conversion in OneHot preprocessor (would cause float64 in eager, then NN-matmul-failure).
Add proper seq-len + state-in construction in eager_tf_policy.py::_compute_gradients().

* LINT.

* eager_tf_policy.py: Only set samples["seq_lens"] if RNN. Otherwise, eager-tracing will throw flattened-dict key-mismatch error.

* Move issue code to examples folder.

Co-authored-by: Eric Liang <ekhliang@gmail.com>

rllib/evaluation/sampler.py

@@ -525,10 +525,11 @@ def _do_policy_eval(tf_sess, to_eval, policies, active_episodes):
             summarize(to_eval)))
 
     for policy_id, eval_data in to_eval.items():
-        rnn_in_cols = _to_column_format([t.rnn_state for t in eval_data])
+        rnn_in = [t.rnn_state for t in eval_data]
         policy = _get_or_raise(policies, policy_id)
         if builder and (policy.compute_actions.__code__ is
                         TFPolicy.compute_actions.__code__):
+            rnn_in_cols = _to_column_format(rnn_in)
             # TODO(ekl): how can we make info batch available to TF code?
             pending_fetches[policy_id] = policy._build_compute_actions(
                 builder, [t.obs for t in eval_data],
@@ -536,6 +537,11 @@ def _do_policy_eval(tf_sess, to_eval, policies, active_episodes):
                 prev_action_batch=[t.prev_action for t in eval_data],
                 prev_reward_batch=[t.prev_reward for t in eval_data])
         else:
+            # TODO(sven): Does this work for LSTM torch?
+            rnn_in_cols = [
+                np.stack([row[i] for row in rnn_in])
+                for i in range(len(rnn_in[0]))
+            ]
             eval_results[policy_id] = policy.compute_actions(
                 [t.obs for t in eval_data],
                 rnn_in_cols,

rllib/examples/custom_keras_cnn_plus_rnn_model.py

@@ -0,0 +1,115 @@
+# Explains/tests Issues:
+# https://github.com/ray-project/ray/issues/6928
+# https://github.com/ray-project/ray/issues/6732
+
+from gym.spaces import Discrete, Box
+import numpy as np
+
+from ray.rllib.agents.ppo import PPOTrainer
+from ray.rllib.examples.random_env import RandomEnv
+from ray.rllib.models import ModelCatalog
+from ray.rllib.models.modelv2 import ModelV2
+from ray.rllib.models.tf.recurrent_tf_modelv2 import RecurrentTFModelV2
+from ray.rllib.utils import try_import_tf
+from ray.rllib.utils.annotations import override
+
+tf = try_import_tf()
+
+cnn_shape = (4, 4, 3)
+
+
+class CustomModel(RecurrentTFModelV2):
+    def __init__(self, obs_space, action_space, num_outputs, model_config,
+                 name):
+        super(CustomModel, self).__init__(obs_space, action_space, num_outputs,
+                                          model_config, name)
+
+        self.cell_size = 16
+        visual_size = cnn_shape[0] * cnn_shape[1] * cnn_shape[2]
+
+        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)
+
+        inputs = tf.keras.layers.Input(
+            shape=(None, visual_size), name="visual_inputs")
+
+        input_visual = inputs
+        input_visual = tf.reshape(
+            input_visual, [-1, cnn_shape[0], cnn_shape[1], cnn_shape[2]])
+        cnn_input = tf.keras.layers.Input(shape=cnn_shape, name="cnn_input")
+
+        cnn_model = tf.keras.applications.mobilenet_v2.MobileNetV2(
+            alpha=1.0,
+            include_top=True,
+            weights=None,
+            input_tensor=cnn_input,
+            pooling=None)
+        vision_out = cnn_model(input_visual)
+        vision_out = tf.reshape(
+            vision_out,
+            [-1, tf.shape(inputs)[1],
+             vision_out.shape.as_list()[-1]])
+
+        lstm_out, state_h, state_c = tf.keras.layers.LSTM(
+            self.cell_size,
+            return_sequences=True,
+            return_state=True,
+            name="lstm")(
+                inputs=vision_out,
+                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=[inputs, 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(RecurrentTFModelV2)
+    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])
+
+
+if __name__ == "__main__":
+    ModelCatalog.register_custom_model("my_model", CustomModel)
+    trainer = PPOTrainer(
+        env=RandomEnv,
+        config={
+            # "eager": True,  # <- should work for both eager or not
+            "model": {
+                "custom_model": "my_model",
+                "max_seq_len": 20,
+            },
+            "vf_share_layers": True,
+            "num_workers": 0,  # no parallelism
+            "env_config": {
+                "action_space": Discrete(2),
+                # Test a simple Tuple observation space.
+                "observation_space": Box(
+                    0.0, 1.0, shape=cnn_shape, dtype=np.float32)
+            }
+        })
+    trainer.train()

rllib/models/preprocessors.py

@@ -143,7 +143,7 @@ class OneHotPreprocessor(Preprocessor):
     @override(Preprocessor)
     def transform(self, observation):
         self.check_shape(observation)
-        arr = np.zeros(self._obs_space.n)
+        arr = np.zeros(self._obs_space.n, dtype=np.float32)
         arr[observation] = 1
         return arr
 

rllib/policy/eager_tf_policy.py

@@ -227,16 +227,18 @@ def build_eager_tf_policy(name,
                     framework="tf",
                 )
 
+            self._state_in = [
+                tf.convert_to_tensor(np.array([s]))
+                for s in self.model.get_initial_state()
+            ]
+
             self.model({
                 SampleBatch.CUR_OBS: tf.convert_to_tensor(
                     np.array([observation_space.sample()])),
                 SampleBatch.PREV_ACTIONS: tf.convert_to_tensor(
                     [_flatten_action(action_space.sample())]),
                 SampleBatch.PREV_REWARDS: tf.convert_to_tensor([0.]),
-            }, [
-                tf.convert_to_tensor(np.array([s]))
-                for s in self.model.get_initial_state()
-            ], tf.convert_to_tensor([1]))
+            }, self._state_in, tf.convert_to_tensor([1]))
 
             if before_loss_init:
                 before_loss_init(self, observation_space, action_space, config)
@@ -303,7 +305,7 @@ def build_eager_tf_policy(name,
             else:
                 n = obs_batch.shape[0]
 
-            seq_lens = tf.ones(n)
+            seq_lens = tf.ones(n, dtype=tf.int32)
             input_dict = {
                 SampleBatch.CUR_OBS: tf.convert_to_tensor(obs_batch),
                 "is_training": tf.constant(False),
@@ -369,6 +371,10 @@ def build_eager_tf_policy(name,
         def num_state_tensors(self):
             return len(self._state_in)
 
+        @override(Policy)
+        def get_initial_state(self):
+            return self.model.get_initial_state()
+
         def get_session(self):
             return None  # None implies eager
 
@@ -404,9 +410,18 @@ def build_eager_tf_policy(name,
             self._is_training = True
 
             with tf.GradientTape(persistent=gradients_fn is not None) as tape:
-                # TODO: set seq len and state in properly
-                self._seq_lens = tf.ones(samples[SampleBatch.CUR_OBS].shape[0])
-                self._state_in = []
+                # TODO: set seq len and state-in properly
+                state_in = []
+                for i in range(self.num_state_tensors()):
+                    state_in.append(samples["state_in_{}".format(i)])
+                self._state_in = state_in
+
+                self._seq_lens = None
+                if len(state_in) > 0:
+                    self._seq_lens = tf.ones(
+                        samples[SampleBatch.CUR_OBS].shape[0], dtype=tf.int32)
+                    samples["seq_lens"] = self._seq_lens
+
                 model_out, _ = self.model(samples, self._state_in,
                                           self._seq_lens)
                 loss = loss_fn(self, self.model, self.dist_class, samples)
@@ -481,16 +496,11 @@ def build_eager_tf_policy(name,
                     SampleBatch.PREV_ACTIONS: dummy_batch[SampleBatch.ACTIONS],
                     SampleBatch.PREV_REWARDS: dummy_batch[SampleBatch.REWARDS],
                 })
-            state_init = self.get_initial_state()
-            state_batches = []
-            for i, h in enumerate(state_init):
-                dummy_batch["state_in_{}".format(i)] = tf.convert_to_tensor(
-                    np.expand_dims(h, 0))
-                dummy_batch["state_out_{}".format(i)] = tf.convert_to_tensor(
-                    np.expand_dims(h, 0))
-                state_batches.append(
-                    tf.convert_to_tensor(np.expand_dims(h, 0)))
-            if state_init:
+            for i, h in enumerate(self._state_in):
+                dummy_batch["state_in_{}".format(i)] = h
+                dummy_batch["state_out_{}".format(i)] = h
+
+            if self._state_in:
                 dummy_batch["seq_lens"] = tf.convert_to_tensor(
                     np.array([1], dtype=np.int32))
 
@@ -508,7 +518,7 @@ def build_eager_tf_policy(name,
             # Execute a forward pass to get self.action_dist etc initialized,
             # and also obtain the extra action fetches
             _, _, fetches = self.compute_actions(
-                dummy_batch[SampleBatch.CUR_OBS], state_batches,
+                dummy_batch[SampleBatch.CUR_OBS], self._state_in,
                 dummy_batch.get(SampleBatch.PREV_ACTIONS),
                 dummy_batch.get(SampleBatch.PREV_REWARDS))
             dummy_batch.update(fetches)

rllib/tests/test_eager_support.py

@@ -39,8 +39,7 @@ class TestEagerSupport(unittest.TestCase):
         check_support("A2C", {"num_workers": 0})
 
     def testA3C(self):
-        # TODO(ekl) trace on is flaky
-        check_support("A3C", {"num_workers": 1}, test_trace=False)
+        check_support("A3C", {"num_workers": 1})
 
     def testPG(self):
         check_support("PG", {"num_workers": 0})