CEMPlanner

class torchrl.modules.CEMPlanner(*args, **kwargs)

CEMPlanner 模組。

參考文獻：The cross-entropy method for optimization, Botev et al. 2013

給定包含初始狀態的 TensorDict 時，此模組將執行 CEM 規劃步驟。 CEM 規劃步驟的執行方式是從具有零均值和單位變異數的高斯分佈中取樣動作。 然後使用取樣的動作在環境中執行 rollout。 然後對 rollout 獲得的累積獎勵進行排名。 我們選擇前 k 個 episode，並使用它們的動作來更新動作分佈的平均值和標準差。 CEM 規劃步驟重複指定的步驟數。

呼叫模組會傳回根據經驗最大化給定規劃範圍的回報的動作

參數:

• env (EnvBase) – 用於執行規劃步驟的環境 (可以是 ModelBasedEnv 或 EnvBase)。

• planning_horizon (int) – 模擬軌跡的長度

• optim_steps (int) – MPC 規劃器使用的最佳化步驟數

• num_candidates (int) – 從高斯分佈中取樣的候選數量。

• top_k (int) – 用於更新高斯分佈的平均值和標準差的前 k 個候選數量。

• reward_key (str, optional) – 用於檢索獎勵的 TensorDict 中的鍵。預設為“reward”。

• action_key (str, optional) – 用於儲存動作的 TensorDict 中的鍵。預設為“action”

範例

>>> from tensordict import TensorDict
>>> from torchrl.data import Composite, Unbounded
>>> from torchrl.envs.model_based import ModelBasedEnvBase
>>> from torchrl.modules import SafeModule
>>> class MyMBEnv(ModelBasedEnvBase):
...     def __init__(self, world_model, device="cpu", dtype=None, batch_size=None):
...         super().__init__(world_model, device=device, dtype=dtype, batch_size=batch_size)
...         self.state_spec = Composite(
...             hidden_observation=Unbounded((4,))
...         )
...         self.observation_spec = Composite(
...             hidden_observation=Unbounded((4,))
...         )
...         self.action_spec = Unbounded((1,))
...         self.reward_spec = Unbounded((1,))
...
...     def _reset(self, tensordict: TensorDict) -> TensorDict:
...         tensordict = TensorDict(
...             {},
...             batch_size=self.batch_size,
...             device=self.device,
...         )
...         tensordict = tensordict.update(
...             self.full_state_spec.rand())
...         tensordict = tensordict.update(
...             self.full_action_spec.rand())
...         tensordict = tensordict.update(
...             self.full_observation_spec.rand())
...         return tensordict
...
>>> from torchrl.modules import MLP, WorldModelWrapper
>>> import torch.nn as nn
>>> world_model = WorldModelWrapper(
...     SafeModule(
...         MLP(out_features=4, activation_class=nn.ReLU, activate_last_layer=True, depth=0),
...         in_keys=["hidden_observation", "action"],
...         out_keys=["hidden_observation"],
...     ),
...     SafeModule(
...         nn.Linear(4, 1),
...         in_keys=["hidden_observation"],
...         out_keys=["reward"],
...     ),
... )
>>> env = MyMBEnv(world_model)
>>> # Build a planner and use it as actor
>>> planner = CEMPlanner(env, 10, 11, 7, 3)
>>> env.rollout(5, planner)
TensorDict(
    fields={
        action: Tensor(shape=torch.Size([5, 1]), device=cpu, dtype=torch.float32, is_shared=False),
        done: Tensor(shape=torch.Size([5, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        hidden_observation: Tensor(shape=torch.Size([5, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        next: TensorDict(
            fields={
                done: Tensor(shape=torch.Size([5, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                hidden_observation: Tensor(shape=torch.Size([5, 4]), device=cpu, dtype=torch.float32, is_shared=False),
                reward: Tensor(shape=torch.Size([5, 1]), device=cpu, dtype=torch.float32, is_shared=False),
                terminated: Tensor(shape=torch.Size([5, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
            batch_size=torch.Size([5]),
            device=cpu,
            is_shared=False),
        terminated: Tensor(shape=torch.Size([5, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
    batch_size=torch.Size([5]),
    device=cpu,
    is_shared=False)

planning(tensordict: TensorDictBase) → Tensor

執行 MPC 規劃步驟。

參數:

td (TensorDict) – 要在其上執行規劃步驟的 TensorDict。