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.

  • 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.

get_reward_operator() TensorDictModule[source]

Returns a reward operator that maps a world state to a reward.

get_transition_model_operator() TensorDictModule[source]

Returns a transition operator that maps either an observation to a world state or a world state to the next world state.


Access comprehensive developer documentation for PyTorch

View Docs


Get in-depth tutorials for beginners and advanced developers

View Tutorials


Find development resources and get your questions answered

View Resources