[RLlib] Support native tf.keras.Model (milestone toward obsoleting ModelV2 class). #14684

sven1977 · 2021-03-15T13:03:15Z

This PR adds support for native keras models in TF- and eager Policies.

This is an experimental API extension that's fully backward compatible (ModelV2s will still work, of course!).

Users can define a pure tf.keras.Model as their model and have the TF/eager policy use it via the "custom_model" config key.
The keras model only has to implement a get_initial_state method (this requirement may be removed as well to support functional model building or using tf.keras.Sequential). All other ModelV2 methods are no longer required.
The keras.Model call method must be implemented for the forward logic.
An example (RNN) model class has been added (examples/models/modelv3.py), as well as a test case confirming that PPO is learning CartPole using this approach.
A new (experimental) config flag "_use_default_native_models" (False by default) has been added to catalog.py. If True, RLlib will use built-in keras default models where possible (e.g. done in the PPO test). Setting this flag to True is save (RLlib will not try to use keras models where this is not supported, e.g. for torch, for VisionNets, etc..).

NOTE: This PR is for TF only. A solution for supporting pure native torch.Modules is to follow.

Follow up PRs will include:

Making all default tf nets native keras (VisionNet + RNN + LSTM-/attention-wrappers).
Support for PyTorch native Modules.
Examples for custom Sequential types, RNNs (how to specify initial states?), trajectory view API, etc..

Why are these changes needed?

Related issue number

Checks

I've run scripts/format.sh to lint the changes in this PR.
I've included any doc changes needed for https://docs.ray.io/en/master/.
I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
Testing Strategy
- Unit tests
- Release tests
- This PR is not tested :(

…dict_via_sample_batch # Conflicts: # rllib/policy/sample_batch.py

…lete_usage_tracking_dict_via_sample_batch

…lete_usage_tracking_dict_via_sample_batch # Conflicts: # rllib/agents/marwil/marwil_tf_policy.py # rllib/agents/marwil/marwil_torch_policy.py # rllib/models/tf/recurrent_net.py # rllib/models/torch/recurrent_net.py # rllib/offline/json_reader.py # rllib/policy/sample_batch.py

…lete_usage_tracking_dict_via_sample_batch # Conflicts: # rllib/agents/ppo/ppo_tf_policy.py # rllib/models/modelv2.py # rllib/policy/dynamic_tf_policy.py # rllib/policy/eager_tf_policy.py # rllib/policy/policy.py # rllib/policy/sample_batch.py # rllib/policy/torch_policy.py # rllib/utils/tracking_dict.py

…lete_usage_tracking_dict_via_sample_batch

…lete_modelv2_class # Conflicts: # rllib/agents/ddpg/ddpg_torch_policy.py # rllib/evaluation/postprocessing.py # rllib/models/modelv2.py # rllib/policy/dynamic_tf_policy.py # rllib/policy/sample_batch.py

…lete_modelv2_class

sven1977 · 2021-04-13T08:00:15Z

rllib/agents/marwil/marwil_tf_policy.py

@@ -67,7 +67,8 @@ def postprocess_advantages(policy,
        # input_dict.
        # 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)
+        input_dict = sample_batch.get_single_step_input_dict(


Moved this helper method to SampleBatch (from ModelV2), where it's more fitting.

sven1977 · 2021-04-13T08:01:25Z

rllib/agents/ppo/ppo_tf_policy.py

        dist_class (Type[ActionDistribution]: The action distr. class.
        train_batch (SampleBatch): The training data.

    Returns:
        Union[TensorType, List[TensorType]]: A single loss tensor or a list
            of loss tensors.
    """
-    logits, state = model.from_batch(train_batch)
+    if isinstance(model, tf.keras.Model):


New model API: native + return 3 values (model-out, state-outs, extra-outs (e.g. vf)).

Assuming ModelV2 will be deprecated soon :)

sven1977 · 2021-04-13T08:02:08Z

rllib/agents/ppo/ppo_tf_policy.py

@@ -86,7 +94,6 @@ def reduce_mean_valid(t):

    if policy.config["use_gae"]:
        prev_value_fn_out = train_batch[SampleBatch.VF_PREDS]
-        value_fn_out = model.value_function()


One reason for the new API is that we can get rid of these async value_function calls.

sven1977 · 2021-04-13T08:02:57Z

rllib/agents/ppo/tests/test_ppo.py

@@ -82,6 +82,9 @@ def test_ppo_compilation_and_lr_schedule(self):
        # Settings in case we use an LSTM.
        config["model"]["lstm_cell_size"] = 10
        config["model"]["max_seq_len"] = 20
+        # Use default-native keras model whenever possible.


Testing, whether this flag works (try using native keras default model).
Currently only works for FCNets.

sven1977 · 2021-04-13T08:03:37Z

rllib/examples/models/modelv3.py

+        values = tf.reshape(self.values(lstm_out), [-1])
+        return logits, [h, c], {SampleBatch.VF_PREDS: values}
+
+    def get_initial_state(self):


We should be able to get rid of this as well in the future via the trajectory view API.

Some users would want different types of initial states (e.g. xavier initalization) right?

sven1977 · 2021-04-13T09:11:35Z

rllib/models/modelv2.py

@@ -333,71 +333,6 @@ def is_time_major(self) -> bool:
        """
        return self.time_major is True

-    # TODO: (sven) Experimental method.


Was moved to SampleBatch instead (makes more sense there).
SampleBatch converts itself into an single-timestep input-dict (which is also a SampleBatch).

sven1977 · 2021-04-13T09:15:24Z

rllib/policy/sample_batch.py

@@ -588,6 +579,71 @@ def data(self):
                old="SampleBatch.data[..]", new="SampleBatch[..]", error=False)
        return self

+    # TODO: (sven) Experimental method.


Moved here from ModelV2 (makes more sense inside SampleBatch).

…lete_modelv2_class

…lete_modelv2_class # Conflicts: # rllib/agents/marwil/marwil_tf_policy.py

…lete_modelv2_class

…lete_modelv2_class # Conflicts: # rllib/policy/torch_policy.py

…lete_modelv2_class # Conflicts: # rllib/agents/trainer.py

michaelzhiluo · 2021-04-27T00:28:19Z

rllib/models/tf/fcnet.py

+        i = 1
+
+        # Create layers 0 to second-last.
+        for size in hiddens[:-1]:


Is this code copy pasted from the prior modelv2 fcnet for tf?

smorad · 2021-05-06T12:33:10Z

@sven1977
This breaks rllib if tensorflow is not installed.

ray/rllib/models/catalog.py

Line 741 in e973b72

if issubclass(model_class, tf.keras.Model):

Should be if if tf and issubclass(model_class, tf.keras.Model):

and

ray/rllib/models/catalog.py

Line 408 in e973b72

if framework not in ["tf", "tf2", "tfe"] or \

Should be if tf and framework not in ["tf", "tf2", "tfe"]

I'd patch this but I've managed to totally mangle my fork/local repo atm.

sven1977 added 28 commits December 23, 2020 17:26

WIP.

7c5f7ce

WIP.

8f74a7b

WIP.

a8cb715

WIP.

43fcfc5

Fix.

8f95df7

Merge branch 'bc_marwil_minor_cleanups' into obsolete_usage_tracking_…

ef80373

…dict_via_sample_batch # Conflicts: # rllib/policy/sample_batch.py

Merge branch 'master' of https://github.com/ray-project/ray into obso…

8dd46c2

…lete_usage_tracking_dict_via_sample_batch

Fixes and LINT.

40a6d8b

Fixes and LINT.

bb728ef

WIP.

4f16efe

wip

371a298

wip

bdb7c93

wip

f03f49d

wip

d4e31de

wip

1988c53

wip

493424a

wip

a74d668

wip

a979ca8

Merge branch 'master' of https://github.com/ray-project/ray into obso…

e93354a

…lete_usage_tracking_dict_via_sample_batch

wip

3fa7e55

wip

deaa7f5

Merge branch 'master' of https://github.com/ray-project/ray into obso…

287ea97

…lete_modelv2_class # Conflicts: # rllib/agents/ddpg/ddpg_torch_policy.py # rllib/evaluation/postprocessing.py # rllib/models/modelv2.py # rllib/policy/dynamic_tf_policy.py # rllib/policy/sample_batch.py

wip.

214332a

Merge branch 'master' of https://github.com/ray-project/ray into obso…

94f8f71

…lete_modelv2_class

fix.

b92c702

wip.

b393f03

sven1977 marked this pull request as ready for review March 23, 2021 13:55

sven1977 requested a review from michaelzhiluo March 23, 2021 13:55

sven1977 added 9 commits April 10, 2021 22:11

wip.

0a5bf55

fix

32f5140

fix and LINT.

131ec8f

Merge branch 'master' of https://github.com/ray-project/ray into obso…

7c11cec

…lete_modelv2_class

fixes.

76cdc92

Merge branch 'master' of https://github.com/ray-project/ray into obso…

b8ea5a1

…lete_modelv2_class

fixes.

301e6f9

fix and LINT.

f097b81

Merge branch 'master' of https://github.com/ray-project/ray into obso…

6ddb1a6

…lete_modelv2_class

sven1977 added the tests-ok The tagger certifies test failures are unrelated and assumes personal liability. label Apr 13, 2021

sven1977 commented Apr 13, 2021

View reviewed changes

sven1977 added 2 commits April 13, 2021 11:24

wip.

86c6349

Merge branch 'master' of https://github.com/ray-project/ray into obso…

e07ed80

…lete_modelv2_class

sven1977 mentioned this pull request Apr 13, 2021

[RLlib] Support native tf.keras.Models (part 2) - Default keras models for Vision/RNN/Attention. #15273

Merged

6 tasks

sven1977 added 5 commits April 15, 2021 14:59

Merge branch 'master' of https://github.com/ray-project/ray into obso…

4677ee5

…lete_modelv2_class

Merge branch 'master' of https://github.com/ray-project/ray into obso…

400092c

…lete_modelv2_class # Conflicts: # rllib/agents/marwil/marwil_tf_policy.py

Merge branch 'master' of https://github.com/ray-project/ray into obso…

469fca6

…lete_modelv2_class

Merge branch 'master' of https://github.com/ray-project/ray into obso…

ab31bf5

…lete_modelv2_class # Conflicts: # rllib/policy/torch_policy.py

Merge branch 'master' of https://github.com/ray-project/ray into obso…

d768a0e

…lete_modelv2_class # Conflicts: # rllib/agents/trainer.py

michaelzhiluo reviewed Apr 27, 2021

View reviewed changes

michaelzhiluo approved these changes Apr 27, 2021

View reviewed changes

sven1977 merged commit bb8a286 into ray-project:master Apr 27, 2021

sven1977 deleted the obsolete_modelv2_class branch June 2, 2023 20:14

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[RLlib] Support native tf.keras.Model (milestone toward obsoleting ModelV2 class). #14684

[RLlib] Support native tf.keras.Model (milestone toward obsoleting ModelV2 class). #14684

sven1977 commented Mar 15, 2021 •

edited

Loading

sven1977 Apr 13, 2021

sven1977 Apr 13, 2021

michaelzhiluo Apr 26, 2021

sven1977 Apr 13, 2021

sven1977 Apr 13, 2021

sven1977 Apr 13, 2021

michaelzhiluo Apr 27, 2021

sven1977 Apr 13, 2021

sven1977 Apr 13, 2021

michaelzhiluo Apr 27, 2021

smorad commented May 6, 2021

[RLlib] Support native tf.keras.Model (milestone toward obsoleting ModelV2 class). #14684

[RLlib] Support native tf.keras.Model (milestone toward obsoleting ModelV2 class). #14684

Conversation

sven1977 commented Mar 15, 2021 • edited Loading

Why are these changes needed?

Related issue number

Checks

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

smorad commented May 6, 2021

sven1977 commented Mar 15, 2021 •

edited

Loading