WorldModelWrapper

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

世界模型包装器。

该模块封装了转移模型和奖励模型。转移模型用于预测一个假想的世界状态。奖励模型用于预测假想转移的奖励。

参数:

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

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

get_reward_operator() → TensorDictModule

返回一个奖励操作器，该操作器将世界状态映射到奖励。

get_transition_model_operator() → TensorDictModule

返回一个转移操作器，该操作器将观测值映射到世界状态，或将当前世界状态映射到下一个世界状态。