VmasWrapper

torchrl.envs.libs.vmas.VmasWrapper(*args, _inplace_update=False, _batch_locked=True, **kwargs)

Vmas environment wrapper.

Examples

>>>  env = VmasWrapper(
...      vmas.make_env(
...          scenario="flocking",
...          num_envs=32,
...          continuous_actions=True,
...          max_steps=200,
...          device="cpu",
...          seed=None,
...          # Scenario kwargs
...          n_agents=5,
...      )
...  )
>>>  print(env.rollout(10))
TensorDict(
    fields={
        action: Tensor(shape=torch.Size([32, 10, 5, 2]), device=cpu, dtype=torch.float32, is_shared=False),
        done: Tensor(shape=torch.Size([32, 10, 5, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        info: TensorDict(
            fields={
                agent_collision_rew: Tensor(shape=torch.Size([32, 10, 5, 1]), device=cpu, dtype=torch.float32, is_shared=False),
                agent_distance_rew: Tensor(shape=torch.Size([32, 10, 5, 1]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([32, 10, 5]),
            device=cpu,
            is_shared=False),
        next: TensorDict(
            fields={
                done: Tensor(shape=torch.Size([32, 10, 5, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                info: TensorDict(
                    fields={
                        agent_collision_rew: Tensor(shape=torch.Size([32, 10, 5, 1]), device=cpu, dtype=torch.float32, is_shared=False),
                        agent_distance_rew: Tensor(shape=torch.Size([32, 10, 5, 1]), device=cpu, dtype=torch.float32, is_shared=False)},
                    batch_size=torch.Size([32, 10, 5]),
                    device=cpu,
                    is_shared=False),
                observation: Tensor(shape=torch.Size([32, 10, 5, 18]), device=cpu, dtype=torch.float32, is_shared=False),
                reward: Tensor(shape=torch.Size([32, 10, 5, 1]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([32, 10]),
            device=cpu,
            is_shared=False),
        observation: Tensor(shape=torch.Size([32, 10, 5, 18]), device=cpu, dtype=torch.float32, is_shared=False),
        reward: Tensor(shape=torch.Size([32, 10, 5, 1]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([32, 10]),
    device=cpu,
    is_shared=False)