[Feature] Vectorized lidar (#124)

* init

* added raycasting with multiple angles

* out commenting the old

* aemnd

* aemnd

* aemnd

* aemnd

* aemnd

* aemnd

* aemnd

---------

Co-authored-by: Zartris <Jonas.le.fevre@gmail.com>

Author: matteobettini

tests/test_lidar.py

@@ -0,0 +1,28 @@
+#  Copyright (c) 2024.
+#  ProrokLab (https://www.proroklab.org/)
+#  All rights reserved.
+
+import torch
+
+from vmas import make_env
+
+
+def test_vectorized_lidar(n_envs=12, n_steps=15):
+    def get_obs(env):
+        rollout_obs = []
+        for _ in range(n_steps):
+            obs, _, _, _ = env.step(env.get_random_actions())
+            obs = torch.stack(obs, dim=-1)
+            rollout_obs.append(obs)
+        return torch.stack(rollout_obs, dim=-1)
+
+    env_vec_lidar = make_env(
+        scenario="pollock", num_envs=n_envs, seed=0, lidar=True, vectorized_lidar=True
+    )
+    obs_vec_lidar = get_obs(env_vec_lidar)
+    env_non_vec_lidar = make_env(
+        scenario="pollock", num_envs=n_envs, seed=0, lidar=True, vectorized_lidar=False
+    )
+    obs_non_vec_lidar = get_obs(env_non_vec_lidar)
+
+    assert torch.allclose(obs_vec_lidar, obs_non_vec_lidar)

vmas/scenarios/debug/pollock.py

@@ -7,6 +7,8 @@
 from vmas import render_interactively
 from vmas.simulator.core import Agent, Box, Landmark, Line, Sphere, World
 from vmas.simulator.scenario import BaseScenario
+
+from vmas.simulator.sensors import Lidar
 from vmas.simulator.utils import Color, ScenarioUtils
 
 
@@ -15,6 +17,8 @@ def make_world(self, batch_dim: int, device: torch.device, **kwargs):
         self.n_agents = kwargs.pop("n_agents", 15)
         self.n_lines = kwargs.pop("n_lines", 15)
         self.n_boxes = kwargs.pop("n_boxes", 15)
+        self.lidar = kwargs.pop("lidar", False)
+        self.vectorized_lidar = kwargs.pop("vectorized_lidar", True)
         ScenarioUtils.check_kwargs_consumed(kwargs)
 
         self.agent_radius = 0.05
@@ -43,6 +47,7 @@ def make_world(self, batch_dim: int, device: torch.device, **kwargs):
                 shape=Sphere(radius=self.agent_radius),
                 u_multiplier=0.7,
                 rotatable=True,
+                sensors=[Lidar(world, n_rays=16, max_range=0.5)] if self.lidar else [],
             )
             world.add_agent(agent)
 
@@ -85,7 +90,11 @@ def reward(self, agent: Agent):
         return torch.zeros(self.world.batch_dim, device=self.world.device)
 
     def observation(self, agent: Agent):
-        return torch.zeros(self.world.batch_dim, 1, device=self.world.device)
+        return (
+            torch.zeros(self.world.batch_dim, 1, device=self.world.device)
+            if not self.lidar
+            else agent.sensors[0].measure(vectorized=self.vectorized_lidar)
+        )
 
 
 if __name__ == "__main__":