WorldModelWrapper

class torchrl.modules.tensordict_module.WorldModelWrapper(*args, **kwargs)

世界模型包装器。

此模块将转换模型和奖励模型包装在一起。转换模型用于预测虚构的世界状态。奖励模型用于预测想象的转换的奖励。

参数：

• transition_model (TensorDictModule) – 生成新世界状态的转换模型。

• reward_model (TensorDictModule) – 奖励模型，读取世界状态并返回奖励。

get_reward_operator() → TensorDictModule

返回一个奖励运算符，将世界状态映射到奖励。

get_transition_model_operator() → TensorDictModule

返回一个转换运算符，将观察结果映射到世界状态，或将世界状态映射到下一个世界状态。