BraxEnv¶

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

使用环境名称构建的 Google Brax 环境包装器。

Brax 提供了一个基于 Jax 的矢量化和可微分仿真框架。TorchRL 的包装器对 jax 到 torch 的转换产生了一些开销，但仍然可以在模拟轨迹之上构建计算图，允许通过回滚进行反向传播。

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

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

参数：

env_name (str) – 要包装的环境名称。必须是 available_envs 的一部分。
categorical_action_encoding (bool, optional) – 如果为 True，则将分类规范转换为 TorchRL 等效规范 (torchrl.data.DiscreteTensorSpec)，否则将使用独热编码 (torchrl.data.OneHotTensorSpec)。默认为 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() 后 env 处于 done 状态。默认为 False。

变量：

available_envs – 可用于构建的环境

示例

>>> from torchrl.envs import BraxEnv
>>> import torch
>>> device = "cuda" if torch.cuda.is_available() else "cpu"
>>> env = BraxEnv("ant", 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，通常会同时执行多个环境。在以下示例中，我们将迭代地测试不同的批次大小并报告短回滚的执行时间

示例

>>> 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'''
... from torchrl.envs import BraxEnv
... env = BraxEnv("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) 设置：[…] 268.46 ms 1 个测量值，10 次运行，1 个线程

128 env.rollout(100) 设置：[…] 433.80 ms 1 个测量值，10 次运行，1 个线程

可以通过回滚进行反向传播并直接优化策略

>>> from torchrl.envs import BraxEnv
>>> from tensordict.nn import TensorDictModule
>>> from torch import nn
>>> import torch
>>>
>>> env = BraxEnv("ant", batch_size=[10], requires_grad=True)
>>> 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)

BraxEnv¶

文档

教程

资源