wehub-resource-sync
32 lines
1.0 KiB
Python
32 lines
1.0 KiB
Python
"""
|
|
[1] Mastering Diverse Domains through World Models - 2023
|
|
D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap
|
|
https://arxiv.org/pdf/2301.04104v1.pdf
|
|
|
|
[2] Mastering Atari with Discrete World Models - 2021
|
|
D. Hafner, T. Lillicrap, M. Norouzi, J. Ba
|
|
https://arxiv.org/pdf/2010.02193.pdf
|
|
"""
|
|
from ray.rllib.core.learner.learner import Learner
|
|
from ray.rllib.utils.annotations import (
|
|
OverrideToImplementCustomLogic_CallToSuperRecommended,
|
|
override,
|
|
)
|
|
|
|
|
|
class DreamerV3Learner(Learner):
|
|
"""DreamerV3 specific Learner class.
|
|
|
|
Only implements the `after_gradient_based_update()` method to define the logic
|
|
for updating the critic EMA-copy after each training step.
|
|
"""
|
|
|
|
@OverrideToImplementCustomLogic_CallToSuperRecommended
|
|
@override(Learner)
|
|
def after_gradient_based_update(self, *, timesteps):
|
|
super().after_gradient_based_update(timesteps=timesteps)
|
|
|
|
# Update EMA weights of the critic.
|
|
for module_id, module in self.module._rl_modules.items():
|
|
module.unwrapped().critic.update_ema()
|