WorldModelWrapper¶
- class torchrl.modules.tensordict_module.WorldModelWrapper(*args, **kwargs)[source]¶
World model wrapper.
This module wraps together a transition model and a reward model. The transition model is used to predict an imaginary world state. The reward model is used to predict the reward of the imagined transition.
- Parameters:
transition_model (TensorDictModule) – a transition model that generates a new world states.
reward_model (TensorDictModule) – a reward model, that reads the world state and returns a reward.