[RLlib] DreamerV3: Catalog enhancements (MLP/CNN encoders/heads completed and unified accross DL frameworks). #33967

sven1977 · 2023-03-31T11:20:09Z

DreamerV3: Catalog enhancements (MLP/CNN encoders/heads completed and unified accross DL frameworks).

MLP/CNN heads/encoders catalog completed
added use_bias option
added use_layernorm option
unified across DL frameworks
more tests (amongst other things: making sure, both torch and tf2 corresponding models have exact same number of trainable and non-trainable params and compute the exact same output values given equal weights and inputs)

Why are these changes needed?

Related issue number

Checks

I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
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 added any new APIs to the API Reference. For example, if I added a
  method in Tune, I've added it in doc/source/tune/api/ under the
  corresponding .rst file.
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 :(

Signed-off-by: sven1977 <[email protected]>

…mer_v3_catalog_enhancements_01

Signed-off-by: sven1977 <[email protected]>

…mer_v3_catalog_enhancements_01

Signed-off-by: sven1977 <[email protected]>

sven1977 · 2023-03-31T15:14:52Z

rllib/BUILD

@@ -1861,6 +1861,27 @@ py_test(
    srcs = ["core/models/tests/test_catalog.py"]
 )

+py_test(


Unified tests between tf and torch to be able to compare exact number of model parameters (assert architectures are the same).

sven1977 · 2023-03-31T15:15:19Z

rllib/core/models/base.py

@@ -20,16 +20,25 @@
 CRITIC: str = "critic"


+def _raise_not_decorated_exception(class_and_method, input_or_output):


Moved this here. This was duplicated code in tf and torch versions.

sven1977 · 2023-03-31T15:15:38Z

rllib/core/models/base.py

@@ -193,6 +202,11 @@ def _forward(self, input_dict: NestedDict, **kwargs) -> NestedDict:
        """
        raise NotImplementedError

+    @abc.abstractmethod


Added this convenience method to the top-level API.

sven1977 · 2023-03-31T15:16:11Z

rllib/core/models/configs.py

    Attributes:
        hidden_layer_dims: The sizes of the hidden layers.
        hidden_layer_activation: The activation function to use after each layer (
            except for the output).
+        hidden_layer_use_layernorm: Whether to insert a LayerNorm functionality


Added to all primitives (MLP and CNN):

option to switch on layernorm'ing in between layers

use bias or not

sven1977 · 2023-03-31T15:17:17Z

rllib/core/models/configs.py

@@ -216,37 +236,47 @@ class CNNEncoderConfig(ModelConfig):
    Attributes:
        input_dims: The input dimension of the network. These must be given in the
            form of `(width, height, channels)`.
-        filter_specifiers: A list of lists, where each element of an inner list
+        cnn_filter_specifiers: A list of lists, where each element of an inner list


Trying to chose more accucate terms for CNN configs.

sven1977 · 2023-03-31T15:17:50Z

rllib/core/models/tests/test_cnn_encoders.py

@@ -0,0 +1,129 @@
+import unittest


Mostly the exact same file as before, but now unified between tf and torch.

sven1977 · 2023-03-31T15:18:01Z

rllib/core/models/tests/test_mlp_encoders.py

@@ -0,0 +1,116 @@
+import unittest


Mostly the exact same file as before, but now unified between tf and torch.

sven1977 · 2023-03-31T15:19:19Z

rllib/core/models/tf/base.py

@@ -0,0 +1,65 @@
+import abc


Moved this from another file. Now this is analogous to the already existing torch/base.py.
Both define the DL-specific base RLlib Model classes TfModel and TorchModel.

sven1977 · 2023-03-31T15:20:05Z

rllib/core/models/tf/encoder.py

@@ -63,12 +71,75 @@ def _forward(self, inputs: NestedDict) -> NestedDict:
        )


+class TfCNNEncoder(TfModel, Encoder):


A new class (tf did not have CNN encoder before).

sven1977 · 2023-03-31T15:21:01Z

rllib/core/models/tf/primitives.py

        self.network = tf.keras.Sequential(layers)

-    def __call__(self, inputs):
+    def call(self, inputs, **kwargs):


I don't think we should touch __call__ directly, but override call instead. At least that's what keras docs say.

Yeah, this turned out to be a problem. Now that we override call, we ran into a naming conflict with keras models, which have their own input_spec. We therefore decided to rename our properties into input_specs (plural) and output_specs, which alleviated this issue. Should be ok now.

Signed-off-by: sven1977 <[email protected]>

sven1977 · 2023-03-31T15:22:57Z

rllib/core/models/torch/encoder.py

        )

    @override(Model)
    def get_input_spec(self) -> Union[Spec, None]:
        return SpecDict(
            {
-                SampleBatch.OBS: TorchTensorSpec("b, h", h=self.config.input_dims[0]),
+                SampleBatch.OBS: TorchTensorSpec("b, d", d=self.config.input_dims[0]),


I renamed these "hanging right-side dims" to d. I feel like h should be used only for image height and LSTM internal h-state.

sven1977 · 2023-03-31T15:23:50Z

rllib/core/models/torch/encoder.py

        layers.append(nn.Flatten())

        # Add a final linear layer to make sure that the outputs have the correct
        # dimensionality.
        layers.append(
            nn.Linear(
-                int(cnn.output_width) * int(cnn.output_height), config.output_dims[0]
+                int(cnn.output_width) * int(cnn.output_height) * int(cnn.output_depth),


We forgot the output depth here.

sven1977 · 2023-03-31T15:23:59Z

rllib/core/models/torch/encoder.py

@@ -104,10 +113,10 @@ def get_input_spec(self) -> Union[Spec, None]:
        return SpecDict(
            {
                SampleBatch.OBS: TorchTensorSpec(
-                    "b, w, h, d",
+                    "b, w, h, c",


ArturNiederfahrenhorst · 2023-04-03T18:25:49Z

rllib/core/models/configs.py

    """
+    def _validate(self, framework: str = "torch"):
+        super()._validate(framework)
+        if self.output_dims is None:


You say in the docstring that this may be None. This conflicts with this check.

fixed.

We are now allowing output_dims=None for any MLP config as for MLPs, simply the last hidden dim could be used.
This is useful for homogenous dense nets, where all layers have the same activation and no special output layer logic is needed.

ArturNiederfahrenhorst · 2023-04-03T18:26:08Z

rllib/core/models/configs.py

+            hidden_layer_use_layernorm=False,
+            output_dims=None,  # maybe None or a 1D tensor
+        )
+        model = config.build()


build should always take in a framework!

great catch! fixed.

ArturNiederfahrenhorst · 2023-04-03T18:26:13Z

rllib/core/models/configs.py

+            output_activation="tanh",
+            use_bias=False,
+        )
+        model = config.build()


ArturNiederfahrenhorst · 2023-04-03T18:49:49Z

rllib/core/models/tf/encoder.py

+            tf.keras.layers.Dense(config.output_dims[0], activation=output_activation),
+        )
+
+        self.net = tf.keras.Sequential(layers)


Nit: Can we unify the constructor a little more? So that the order of things and the comments are the same between Tf and Torch?

@kouroshHakha I've been thinking about our specs and think we should pull the functionality of TfTensorSpec and TorchTensorSpec into TensorSpec and give TensorSpec a framework kwarg.

If the kwarg is None, simply don't enforce tensor framework and check based on the incoming tensor.
That way we could unify our specs..
-> Saves many LOCs
-> Reduce burden auf maintanance and always checking for equality between frameworks

Especially when rolling this out over RLLib and writing many new RLModules and possibly Models over time, the sepcs would become less of a burden.
Wdyt?

This would get us another step close to having 99% of RLModule and Model code in base clases and having the framework specific classes only be separated by an attribute self.framework.
We are almost there for the PPORLModule.

Love the idea. Let's do this!

unified c'tors a little more, comments, structure, etc..

ArturNiederfahrenhorst · 2023-04-03T19:49:02Z

rllib/core/models/torch/primitives.py

    ):
-        """Initialize a TorchCNN object.


Can we unify TorchCNN and TfCNN docstrings?
What is now seen here as "Attributes" should be "Args", right?
I see similar stuff happening with TorchMLP!
Most of what is now listed in the class docstrings as attributes are not actually attributes but still args.

completely unified these now.

you are right, the primitives should have a Args list, not Attributes.
Fixed.

ArturNiederfahrenhorst

Thanks for the massive amounts of small cleanups here!

Signed-off-by: sven1977 <[email protected]>

kouroshHakha

Just a few nit and a question?

kouroshHakha · 2023-04-04T23:31:44Z

rllib/core/models/utils.py

Should we move this to test_utils.py?

kouroshHakha · 2023-04-04T23:32:28Z

rllib/core/models/utils.py

+                else:
+                    inputs[key] = None
+        else:
+            inputs = model.input_specs.fill(self.random_fill_input_value)


haha, I kinda forgot about this fill thing. What a nice use-case man :)

kouroshHakha · 2023-04-04T23:35:53Z

rllib/core/models/utils.py

+        # Bring model into a reproducible, comparable state (so we can compare
+        # computations across frameworks). Use only a value-sequence of len=1 here
+        # as it could possibly be that the layers are stored in different order
+        # across the different frameworks.


if we can't reliably match the order across frameworks, does it make sense for us to support a sequence of values in _set_to_dummy_weights?

Yeah, good point. Probably not. But one might also want to just make a network repeatable compute the same outputs across network instantiations (over time), not necessarily across different frameworks.

kouroshHakha · 2023-04-04T23:37:33Z

rllib/core/models/utils.py

+        main_key = next(iter(self.models.keys()))
+        # Compare number of trainable and non-trainable params between all
+        # frameworks.
+        for c in self.param_counts.values():


can we separate out param count and output checker into two different functions with their own control over accepted tolerance?

kouroshHakha · 2023-04-04T23:46:03Z

rllib/core/models/tf/mlp.py

@@ -27,11 +27,11 @@ def __init__(self, config: MLPHeadConfig) -> None:
        )

    @override(Model)
-    def get_input_spec(self) -> Union[Spec, None]:
+    def get_input_specs(self) -> Union[Spec, None]:


Nit: Optional[Spec]

kouroshHakha · 2023-04-04T23:46:14Z

rllib/core/models/tf/mlp.py

        )

    @override(Model)
-    def get_input_spec(self) -> Union[Spec, None]:
+    def get_input_specs(self) -> Union[Spec, None]:


kouroshHakha · 2023-04-04T23:48:55Z

rllib/core/models/torch/base.py

@@ -104,3 +106,12 @@ def get_num_parameters(self) -> Tuple[int, int]:
            num_trainable_params,
            num_all_params - num_trainable_params,
        )
+
+    @override(Model)
+    def _set_to_dummy_weights(self, value_sequence=(-0.02, -0.01, 0.01, 0.02)):


What is the mechanism to ensure the order of parameters are the same between frameworks?

I'm not sure either. For torch, it's neither the order, in which you define Parameter properties in the ctor, nor alphabetical. Didn't have time to investigate.

but using the same value for all trainable and non-trainable parameters should result in the same behavior between tf and torch. We are relying on this assumption, right?

kouroshHakha · 2023-04-04T23:49:32Z

rllib/core/models/torch/mlp.py

@@ -28,11 +28,11 @@ def __init__(self, config: MLPHeadConfig) -> None:
        )

    @override(Model)
-    def get_input_spec(self) -> Union[Spec, None]:
+    def get_input_specs(self) -> Union[Spec, None]:


kouroshHakha · 2023-04-04T23:49:47Z

rllib/core/models/torch/mlp.py

        )

        self.log_std = torch.nn.Parameter(
            torch.as_tensor([0.0] * self._half_output_dim)
        )

    @override(Model)
-    def get_input_spec(self) -> Union[Spec, None]:
+    def get_input_specs(self) -> Union[Spec, None]:


same: Optional[Spec]

kouroshHakha · 2023-04-04T23:50:54Z

Let's merged upon addressing the questions.

kouroshHakha · 2023-04-04T23:52:19Z

Also tests are failing.

…mer_v3_catalog_enhancements_01

Signed-off-by: sven1977 <[email protected]>

…mer_v3_catalog_enhancements_01

Signed-off-by: sven1977 <[email protected]>

…eted and unified accross DL frameworks). (ray-project#33967) Signed-off-by: elliottower <[email protected]>

…eted and unified accross DL frameworks). (ray-project#33967) Signed-off-by: Jack He <[email protected]>

sven1977 added 7 commits March 28, 2023 16:42

wip

36a4c65

Signed-off-by: sven1977 <[email protected]>

wip

fd432e1

Signed-off-by: sven1977 <[email protected]>

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

eb33ba0

…mer_v3_catalog_enhancements_01

wip

36346cf

Signed-off-by: sven1977 <[email protected]>

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

5995906

…mer_v3_catalog_enhancements_01

wip

950eb81

Signed-off-by: sven1977 <[email protected]>

wip

392b5f6

Signed-off-by: sven1977 <[email protected]>

sven1977 requested review from gjoliver, avnishn, ArturNiederfahrenhorst, smorad, maxpumperla, kouroshHakha and krfricke as code owners March 31, 2023 11:20

sven1977 assigned ArturNiederfahrenhorst Mar 31, 2023

sven1977 mentioned this pull request Mar 31, 2023

[RLlib] DreamerV3: Add Conv2d-transpose support to new model Catalog. #33969

Merged

8 tasks

sven1977 commented Mar 31, 2023

View reviewed changes

fix

46b3c4f

Signed-off-by: sven1977 <[email protected]>

sven1977 commented Mar 31, 2023

View reviewed changes

ArturNiederfahrenhorst reviewed Apr 3, 2023

View reviewed changes

sven1977 added 4 commits April 4, 2023 10:33

merge

3de4cd1

Signed-off-by: sven1977 <[email protected]>

wip

28d0c5e

Signed-off-by: sven1977 <[email protected]>

wip

a9c2897

Signed-off-by: sven1977 <[email protected]>

wip

16101b8

Signed-off-by: sven1977 <[email protected]>

kouroshHakha approved these changes Apr 4, 2023

View reviewed changes

sven1977 added 7 commits April 6, 2023 13:41

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

f16fb63

…mer_v3_catalog_enhancements_01

wip

f349292

Signed-off-by: sven1977 <[email protected]>

LINT

5d6f194

Signed-off-by: sven1977 <[email protected]>

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

0288d21

…mer_v3_catalog_enhancements_01

LINT

6bc283a

Signed-off-by: sven1977 <[email protected]>

LINT

c0b04fa

Signed-off-by: sven1977 <[email protected]>

wip

6506e15

Signed-off-by: sven1977 <[email protected]>

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

sven1977 added 3 commits April 6, 2023 15:38

wip

c6322b0

Signed-off-by: sven1977 <[email protected]>

fix

e59f467

Signed-off-by: sven1977 <[email protected]>

LINT and docs fixes

bfbb3e0

Signed-off-by: sven1977 <[email protected]>

sven1977 requested a review from a team as a code owner April 7, 2023 10:06

sven1977 merged commit 3f6e084 into ray-project:master Apr 7, 2023

elliottower pushed a commit to elliottower/ray that referenced this pull request Apr 22, 2023

[RLlib] DreamerV3: Catalog enhancements (MLP/CNN encoders/heads compl…

279a7d6

…eted and unified accross DL frameworks). (ray-project#33967) Signed-off-by: elliottower <[email protected]>

ProjectsByJackHe pushed a commit to ProjectsByJackHe/ray that referenced this pull request May 4, 2023

[RLlib] DreamerV3: Catalog enhancements (MLP/CNN encoders/heads compl…

8d6476e

…eted and unified accross DL frameworks). (ray-project#33967) Signed-off-by: Jack He <[email protected]>

sven1977 deleted the dreamer_v3_catalog_enhancements_01 branch May 5, 2023 20:04

		@@ -20,16 +20,25 @@
		CRITIC: str = "critic"


		def _raise_not_decorated_exception(class_and_method, input_or_output):

		@@ -63,12 +71,75 @@ def _forward(self, inputs: NestedDict) -> NestedDict:
		)


		class TfCNNEncoder(TfModel, Encoder):

[RLlib] DreamerV3: Catalog enhancements (MLP/CNN encoders/heads completed and unified accross DL frameworks). #33967

[RLlib] DreamerV3: Catalog enhancements (MLP/CNN encoders/heads completed and unified accross DL frameworks). #33967

Conversation

sven1977 commented Mar 31, 2023 • 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

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

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

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ArturNiederfahrenhorst left a comment

Choose a reason for hiding this comment

kouroshHakha left a 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

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

kouroshHakha commented Apr 4, 2023

kouroshHakha commented Apr 4, 2023

sven1977 commented Mar 31, 2023 •

edited

Loading