BraxWrapper¶

torchrl.envs.BraxWrapper(*args, **kwargs)[source]¶

Google Brax 環境包裝器。

Brax 提供了一個基於 Jax 的向量化和可微分模擬框架。TorchRL 的包裝器會產生一些開銷用於 jax-to-torch 轉換，但計算圖仍然可以構建在模擬軌跡之上，從而允許透過 rollouts 進行反向傳播。

GitHub: https://github.com/google/brax

論文: https://arxiv.org/abs/2106.13281

引數:

env (brax.envs.base.PipelineEnv) – 要包裝的環境。
categorical_action_encoding (bool, optional) – 如果為 True，則分類規範將轉換為 TorchRL 等效項 (torchrl.data.Categorical)，否則將使用獨熱編碼 (torchrl.data.OneHot)。預設為 False。
cache_clear_frequency (int, optional) – 在使用 requires_grad=True 時，每 N 步自動清除 JAX 的內部快取以防止記憶體洩漏。預設為 False（停用自動快取清除）。

關鍵字引數:

from_pixels (bool, optional) – 尚不支援。
frame_skip (int, optional) – 如果提供，表示相同的動作需要重複多少步。返回的觀察值將是序列的最後一個觀察值，而獎勵將是跨步驟的獎勵總和。
device (torch.device, optional) – 如果提供，資料將被轉換到的裝置。預設為 torch.device("cpu")。
batch_size (torch.Size, optional) – 環境的批次大小。在 brax 中，這表示向量化環境的數量。預設為 torch.Size([])。
allow_done_after_reset (bool, optional) – 如果為 True，則允許在呼叫 reset() 後立即將環境設定為 done。預設為 False。

變數:

available_envs – 可用於構建的環境

示例

>>> import brax.envs
>>> from torchrl.envs import BraxWrapper
>>> import torch
>>> device = "cuda" if torch.cuda.is_available() else "cpu"
>>> base_env = brax.envs.get_environment("ant")
>>> env = BraxWrapper(base_env, device=device)
>>> env.set_seed(0)
>>> td = env.reset()
>>> td["action"] = env.action_spec.rand()
>>> td = env.step(td)
>>> print(td)
TensorDict(
    fields={
        action: Tensor(torch.Size([8]), dtype=torch.float32),
        done: Tensor(torch.Size([1]), dtype=torch.bool),
        next: TensorDict(
            fields={
                observation: Tensor(torch.Size([87]), dtype=torch.float32)},
            batch_size=torch.Size([]),
            device=cpu,
            is_shared=False),
        observation: Tensor(torch.Size([87]), dtype=torch.float32),
        reward: Tensor(torch.Size([1]), dtype=torch.float32),
        state: TensorDict(...)},
    batch_size=torch.Size([]),
    device=cpu,
    is_shared=False)
>>> print(env.available_envs)
['acrobot', 'ant', 'fast', 'fetch', ...]

為了利用 Brax 的優勢，通常需要同時執行多個環境。在以下示例中，我們迭代地測試不同的批次大小，並報告短 rollout 的執行時間

示例

>>> import torch
>>> from torch.utils.benchmark import Timer
>>> device = "cuda" if torch.cuda.is_available() else "cpu"
>>> for batch_size in [4, 16, 128]:
...     timer = Timer('''
... env.rollout(100)
... ''',
...     setup=f'''
... import brax.envs
... from torchrl.envs import BraxWrapper
... env = BraxWrapper(brax.envs.get_environment("ant"), batch_size=[{batch_size}], device="{device}")
... env.set_seed(0)
... env.rollout(2)
... ''')
...     print(batch_size, timer.timeit(10))
4
env.rollout(100)
setup: [...]
310.00 ms
1 measurement, 10 runs , 1 thread

16 env.rollout(100) setup: […] 268.46 毫秒 1 次測量，10 次執行，1 個執行緒

128 env.rollout(100) setup: […] 433.80 毫秒 1 次測量，10 次執行，1 個執行緒

可以對 rollout 進行反向傳播並直接最佳化策略

>>> import brax.envs
>>> from torchrl.envs import BraxWrapper
>>> from tensordict.nn import TensorDictModule
>>> from torch import nn
>>> import torch
>>>
>>> env = BraxWrapper(brax.envs.get_environment("ant"), batch_size=[10], requires_grad=True, cache_clear_frequency=100)
>>> env.set_seed(0)
>>> torch.manual_seed(0)
>>> policy = TensorDictModule(nn.Linear(27, 8), in_keys=["observation"], out_keys=["action"])
>>>
>>> td = env.rollout(10, policy)
>>>
>>> td["next", "reward"].mean().backward(retain_graph=True)
>>> print(policy.module.weight.grad.norm())
tensor(213.8605)

BraxWrapper¶

文件

教程

資源