Ray rllib DreamerV3 incompatible with new API?

[rschiewer]

What happened + What you expected to happen

I've tried to use the dreamerV3 agent from rllib for one of my projects and find it quite challenging to configure it correctly so that it would not throw an exception. I apologize in advance if this is entirely a problem on my end, but I have reasons to believe that this is not the case. To simplify things, I tried the frozenlake_2x2 tuned example found here
and it results in the same exception I encountered in my code as well:

File "/opt/homebrew/anaconda3/envs/tf/lib/python3.11/site-packages/ray/rllib/algorithms/algorithm.py", line 528, in __init__ super().__init__( File "/opt/homebrew/anaconda3/envs/tf/lib/python3.11/site-packages/ray/tune/trainable/trainable.py", line 157, in __init__ self.setup(copy.deepcopy(self.config)) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/opt/homebrew/anaconda3/envs/tf/lib/python3.11/site-packages/ray/rllib/algorithms/dreamerv3/dreamerv3.py", line 488, in setup super().setup(config) File "/opt/homebrew/anaconda3/envs/tf/lib/python3.11/site-packages/ray/rllib/algorithms/algorithm.py", line 748, in setup self.learner_group = self.config.build_learner_group( ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/opt/homebrew/anaconda3/envs/tf/lib/python3.11/site-packages/ray/rllib/algorithms/algorithm_config.py", line 1249, in build_learner_group learner_group = LearnerGroup(config=self.copy(), module_spec=rl_module_spec) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/opt/homebrew/anaconda3/envs/tf/lib/python3.11/site-packages/ray/rllib/core/learner/learner_group.py", line 129, in __init__ self._learner.build() File "/opt/homebrew/anaconda3/envs/tf/lib/python3.11/site-packages/ray/rllib/core/learner/tf/tf_learner.py", line 275, in build super().build() File "/opt/homebrew/anaconda3/envs/tf/lib/python3.11/site-packages/ray/rllib/core/learner/learner.py", line 310, in build self._module = self._make_module() ^^^^^^^^^^^^^^^^^^^ File "/opt/homebrew/anaconda3/envs/tf/lib/python3.11/site-packages/ray/rllib/core/learner/learner.py", line 1586, in _make_module module = self._module_spec.build() ^^^^^^^^^^^^^^^^^^^^^^^^^ File "/opt/homebrew/anaconda3/envs/tf/lib/python3.11/site-packages/ray/rllib/core/rl_module/multi_rl_module.py", line 611, in build module = self.multi_rl_module_class( ^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/opt/homebrew/anaconda3/envs/tf/lib/python3.11/site-packages/ray/rllib/core/rl_module/multi_rl_module.py", line 127, in __init__ super().__init__( File "/opt/homebrew/anaconda3/envs/tf/lib/python3.11/site-packages/ray/rllib/core/rl_module/rl_module.py", line 467, in __init__ self.setup() File "/opt/homebrew/anaconda3/envs/tf/lib/python3.11/site-packages/ray/rllib/core/rl_module/multi_rl_module.py", line 145, in setup self._rl_modules[module_id] = rl_module_spec.build() ^^^^^^^^^^^^^^^^^^^^^^ File "/opt/homebrew/anaconda3/envs/tf/lib/python3.11/site-packages/ray/rllib/core/rl_module/rl_module.py", line 113, in build module = self.module_class(module_config) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/opt/homebrew/anaconda3/envs/tf/lib/python3.11/site-packages/ray/rllib/core/rl_module/tf/tf_rl_module.py", line 25, in __init__ RLModule.__init__(self, *args, **kwargs) File "/opt/homebrew/anaconda3/envs/tf/lib/python3.11/site-packages/ray/rllib/core/rl_module/rl_module.py", line 408, in __init__ deprecation_warning( File "/opt/homebrew/anaconda3/envs/tf/lib/python3.11/site-packages/ray/rllib/utils/deprecation.py", line 48, in deprecation_warning raise ValueError(msg) ValueError: RLModule(config=[RLModuleConfig]) has been deprecated. Use RLModule(observation_space=.., action_space=.., inference_only=.., learner_only=.., model_config=..) instead.

I've tried to adapt the code to the new rl module API, which I'm admittedly not yet familiar with. I came up with this code:

import ray
from ray.rllib.algorithms.dreamerv3.dreamerv3_rl_module import DreamerV3RLModule
from ray.rllib.algorithms.dreamerv3 import DreamerV3Config
from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
from ray.rllib.core.rl_module.rl_module import RLModuleSpe
ray.init(ignore_reinit_error=True)

config = (
    DreamerV3Config()
    .environment('CartPole-v1')
    .training(
        train_batch_size_per_learner=2000,
        lr=0.0004,
    )
    .framework("tf2")  # set framework explicitly
    .rl_module(
        model_config=DefaultModelConfig(),
        rl_module_spec=RLModuleSpec(
            module_class=DreamerV3RLModule
        )
    )

)

# Build the algorithm.
algo = config.build()

# Training loop.
for i in range(100):
    pprint(algo.train())
ray.shutdown()

This yields the exception:

File "/Users/username/projects/dreamer-cscg/train_dreamer.py", line 98, in <module> main() File "/Users/username/projects/dreamer-cscg/train_dreamer.py", line 89, in main algo = config.build() ^^^^^^^^^^^^^^ File "/opt/homebrew/anaconda3/envs/tf/lib/python3.11/site-packages/ray/rllib/utils/deprecation.py", line 128, in _ctor return obj(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^ File "/opt/homebrew/anaconda3/envs/tf/lib/python3.11/site-packages/ray/rllib/algorithms/algorithm_config.py", line 5417, in build return self.build_algo(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/opt/homebrew/anaconda3/envs/tf/lib/python3.11/site-packages/ray/rllib/algorithms/algorithm_config.py", line 958, in build_algo return algo_class( ^^^^^^^^^^^ File "/opt/homebrew/anaconda3/envs/tf/lib/python3.11/site-packages/ray/rllib/algorithms/algorithm.py", line 528, in __init__ super().__init__( File "/opt/homebrew/anaconda3/envs/tf/lib/python3.11/site-packages/ray/tune/trainable/trainable.py", line 157, in __init__ self.setup(copy.deepcopy(self.config)) File "/opt/homebrew/anaconda3/envs/tf/lib/python3.11/site-packages/ray/rllib/algorithms/dreamerv3/dreamerv3.py", line 488, in setup super().setup(config) File "/opt/homebrew/anaconda3/envs/tf/lib/python3.11/site-packages/ray/rllib/algorithms/algorithm.py", line 748, in setup self.learner_group = self.config.build_learner_group( ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/opt/homebrew/anaconda3/envs/tf/lib/python3.11/site-packages/ray/rllib/algorithms/algorithm_config.py", line 1249, in build_learner_group learner_group = LearnerGroup(config=self.copy(), module_spec=rl_module_spec) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/opt/homebrew/anaconda3/envs/tf/lib/python3.11/site-packages/ray/rllib/core/learner/learner_group.py", line 129, in __init__ self._learner.build() File "/opt/homebrew/anaconda3/envs/tf/lib/python3.11/site-packages/ray/rllib/core/learner/tf/tf_learner.py", line 275, in build super().build() File "/opt/homebrew/anaconda3/envs/tf/lib/python3.11/site-packages/ray/rllib/core/learner/learner.py", line 310, in build self._module = self._make_module() ^^^^^^^^^^^^^^^^^^^ File "/opt/homebrew/anaconda3/envs/tf/lib/python3.11/site-packages/ray/rllib/core/learner/learner.py", line 1586, in _make_module module = self._module_spec.build() ^^^^^^^^^^^^^^^^^^^^^^^^^ File "/opt/homebrew/anaconda3/envs/tf/lib/python3.11/site-packages/ray/rllib/core/rl_module/multi_rl_module.py", line 611, in build module = self.multi_rl_module_class( ^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/opt/homebrew/anaconda3/envs/tf/lib/python3.11/site-packages/ray/rllib/core/rl_module/multi_rl_module.py", line 127, in __init__ super().__init__( File "/opt/homebrew/anaconda3/envs/tf/lib/python3.11/site-packages/ray/rllib/core/rl_module/rl_module.py", line 467, in __init__ self.setup() File "/opt/homebrew/anaconda3/envs/tf/lib/python3.11/site-packages/ray/rllib/core/rl_module/multi_rl_module.py", line 145, in setup self._rl_modules[module_id] = rl_module_spec.build() ^^^^^^^^^^^^^^^^^^^^^^ File "/opt/homebrew/anaconda3/envs/tf/lib/python3.11/site-packages/ray/rllib/core/rl_module/rl_module.py", line 102, in build module = self.module_class( ^^^^^^^^^^^^^^^^^^ File "/opt/homebrew/anaconda3/envs/tf/lib/python3.11/site-packages/ray/rllib/core/rl_module/rl_module.py", line 467, in __init__ self.setup() File "/opt/homebrew/anaconda3/envs/tf/lib/python3.11/site-packages/ray/rllib/algorithms/dreamerv3/dreamerv3_rl_module.py", line 50, in setup self.encoder = self.catalog.build_encoder(framework=self.framework) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/opt/homebrew/anaconda3/envs/tf/lib/python3.11/site-packages/ray/rllib/algorithms/dreamerv3/dreamerv3_catalog.py", line 45, in build_encoder raise NotImplementedError NotImplementedError

Versions / Dependencies

MacOS Sequoia 15.3.1
Python: 3.11

Package Version

---

absl-py 2.1.0
aiohappyeyeballs 2.4.4
aiohttp 3.10.5
aiohttp-cors 0.7.0
aiosignal 1.2.0
annotated-types 0.7.0
astunparse 1.6.3
attrs 24.3.0
blinker 1.9.0
Box2D 2.3.10
Brotli 1.0.9
cachetools 5.5.1
certifi 2025.1.31
cffi 1.17.1
charset-normalizer 3.3.2
click 8.1.7
cloudpickle 3.1.1
colorful 0.5.6
cryptography 44.0.1
decorator 5.2.1
distlib 0.3.9
dm-tree 0.1.9
Farama-Notifications 0.0.4
filelock 3.17.0
flatbuffers 24.3.25
frozenlist 1.5.0
fsspec 2025.3.0
gast 0.4.0
google-api-core 2.24.2
google-auth 2.38.0
google-auth-oauthlib 0.5.2
google-pasta 0.2.0
googleapis-common-protos 1.69.1
grpcio 1.62.2
gymnasium 1.0.0
h5py 3.12.1
idna 3.7
imageio 2.37.0
jinja2 None
jsonschema 4.23.0
jsonschema-specifications 2024.10.1
keras 3.6.0
Keras-Preprocessing 1.1.2
lazy_loader 0.4
libclang 18.1.1
lz4 4.4.3
Markdown 3.4.1
markdown-it-py 2.2.0
MarkupSafe 3.0.2
mdurl 0.1.0
minigrid 3.0.0
ml_dtypes 0.5.1
mpmath 1.3.0
msgpack 1.1.0
multidict 6.1.0
namex 0.0.7
networkx 3.4.2
numpy 2.1.3
oauthlib 3.2.2
opencensus 0.11.4
opencensus-context 0.1.3
opt-einsum 3.3.0
optree 0.14.1
ormsgpack 1.7.0
packaging 24.2
pandas 2.2.3
pillow 11.1.0
pip 25.0.1
platformdirs 4.3.6
prometheus_client 0.21.1
propcache 0.2.0
proto-plus 1.26.1
protobuf 4.25.3
py-spy 0.4.0
pyarrow 19.0.1
pyasn1 0.4.8
pyasn1-modules 0.2.8
pycparser 2.21
pydantic 2.10.6
pydantic_core 2.27.2
pygame 2.6.1
Pygments 2.15.1
PyJWT 2.10.1
pyOpenSSL 25.0.0
PySocks 1.7.1
python-dateutil 2.9.0.post0
pytz 2025.1
PyYAML 6.0.2
ray 2.43.0
referencing 0.36.2
requests 2.32.3
requests-oauthlib 2.0.0
rich 13.9.4
rpds 5.1.0
rpds-py 0.23.1
rsa 4.7.2
scikit-image 0.25.2
scipy 1.15.1
setuptools 75.8.0
shellingham 1.5.4
six 1.16.0
smart-open 7.1.0
svgwrite 1.4.3
swig 4.3.0
sympy 1.13.1
tensorboard 2.19.0
tensorboard_data_server 0.7.0
tensorboard-plugin-wit 1.8.1
tensorboardX 2.6.2.2
tensorflow 2.19.0
tensorflow_estimator 2.15.0
tensorflow-io-gcs-filesystem 0.37.1
tensorflow-probability 0.25.0
termcolor 2.1.0
tf_keras 2.19.0
tifffile 2025.2.18
torch 2.6.0
torchaudio 2.6.0
torchvision 0.21.0
tqdm 4.67.1
typer 0.15.2
typing_extensions 4.12.2
tzdata 2025.1
urllib3 2.3.0
virtualenv 20.29.3
Werkzeug 3.1.3
wheel 0.35.1
wrapt 1.17.0
yarl 1.18.0

Reproduction script

Run the tuned example found here
as instructed inside the script itself. This means:

1. Download run_regression_tests.py
2. Download the example script
3. run with python run_regression_tests.py --dir <absolute_path_to_frozenlake_2x2.py>

Issue Severity

DreamerV3 is not usable for me at the moment.

[walwad2020]

O my god, I just encountered the same problem and I need it for my master's...
Maybe it's time to try out other implementation like ALF.
Sorry ray gods I just can't sift through the documentation anymore