BraxWrapper

torchrl.envs.libs.brax.BraxWrapper(*args, _inplace_update=False, _batch_locked=True, **kwargs)

Google Brax environment wrapper.

Examples

>>> env = brax.envs.get_environment("ant")
>>> env = BraxWrapper(env)
>>> 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', ...]