BraxWrapper¶
- torchrl.envs.BraxWrapper(*args, **kwargs)[源代码]¶
Google Brax 环境包装器。
Brax 提供了一个基于 Jax 的向量化和可微分的模拟框架。TorchRL 的包装器为 jax 到 torch 的转换带来了一些开销,但计算图仍然可以建立在模拟轨迹之上,从而允许通过 rollout 进行反向传播。
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
),否则将使用 one-hot 编码 (torchrl.data.OneHot
)。默认为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()
后,可以容忍 envs 变为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) >>> 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)