continuation of #122

What we vectorize

Let me recap the levels of looping:

1. there are multiple agents
2. each agent can have multiple lidars
3. each lidar sees multtiple entities
4. at multiple angles
5. in multiple environments

And here is what we vectorize:

1. I wouldn't vectorize this as of now as not all agents have lidars
2. i also would not vectorize this as the user might want to read out different lidars at different times (although I would consider it in the future)
3. in this PR
4. in this PR
5. vmas does this already

Benchmarking

pollock scenario

use_vmas_env(
        scenario_name="pollock",
        render=False,
        save_render=False,
        random_action=True,
        continuous_actions=False,
        n_agents=15,
        n_lines=15,
        n_boxes=15,
        lidar=True,
        n_steps=100,
        num_envs=32,
 )

before PR:

M1 cpu: 134.55210709571838s
CUDA: 571.7651047706604s

after PR:

M1 cpu: 22.2061710357666s (0.16x the previous time)
CUDA: 91.27248549461365s (0.15 x the previous time)

benefit will increase in the number of steps and envs

🚀