WorldModelWrapper

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

世界模型包装器。

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

参数:

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

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

get_reward_operator() → TensorDictModule

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

get_transition_model_operator() → TensorDictModule

返回一个转换运算符，该运算符将观测值映射到世界状态，或将世界状态映射到下一个世界状态。