From 5ac5ac9560b1cfa182ed39fcab9931bccaf45bbf Mon Sep 17 00:00:00 2001 From: Sven Mika Date: Thu, 6 Feb 2020 18:44:08 +0100 Subject: [PATCH] [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 --- rllib/evaluation/sampler.py | 8 +- .../custom_keras_cnn_plus_rnn_model.py | 115 ++++++++++++++++++ rllib/models/preprocessors.py | 2 +- rllib/policy/eager_tf_policy.py | 48 +++++--- rllib/tests/test_eager_support.py | 3 +- 5 files changed, 153 insertions(+), 23 deletions(-) create mode 100644 rllib/examples/custom_keras_cnn_plus_rnn_model.py diff --git a/rllib/evaluation/sampler.py b/rllib/evaluation/sampler.py index f03d037e7..41de80405 100644 --- a/rllib/evaluation/sampler.py +++ b/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, diff --git a/rllib/examples/custom_keras_cnn_plus_rnn_model.py b/rllib/examples/custom_keras_cnn_plus_rnn_model.py new file mode 100644 index 000000000..35538c2a3 --- /dev/null +++ b/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() diff --git a/rllib/models/preprocessors.py b/rllib/models/preprocessors.py index 88b253269..162a81f56 100644 --- a/rllib/models/preprocessors.py +++ b/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 diff --git a/rllib/policy/eager_tf_policy.py b/rllib/policy/eager_tf_policy.py index dfbe8c1bb..2a213e8f3 100644 --- a/rllib/policy/eager_tf_policy.py +++ b/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) diff --git a/rllib/tests/test_eager_support.py b/rllib/tests/test_eager_support.py index 143d6970e..348b0461c 100644 --- a/rllib/tests/test_eager_support.py +++ b/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})