[Rendering] More colors and utils to render agent indices and rotation

Author: matteobettini

vmas/scenarios/debug/diff_drive.py

@@ -11,7 +11,7 @@
 from vmas.simulator.dynamics.diff_drive import DiffDrive
 from vmas.simulator.dynamics.holonomic_with_rot import HolonomicWithRotation
 from vmas.simulator.scenario import BaseScenario
-from vmas.simulator.utils import Color, ScenarioUtils
+from vmas.simulator.utils import ScenarioUtils
 
 if typing.TYPE_CHECKING:
     from vmas.simulator.rendering import Geom
@@ -88,24 +88,14 @@ def observation(self, agent: Agent):
         )
 
     def extra_render(self, env_index: int = 0) -> "List[Geom]":
-        from vmas.simulator import rendering
 
         geoms: List[Geom] = []
 
         # Agent rotation
         for agent in self.world.agents:
-            color = Color.BLACK.value
-            line = rendering.Line(
-                (0, 0),
-                (0.1, 0),
-                width=1,
+            geoms.append(
+                ScenarioUtils.plot_entity_rotation(agent, env_index, length=0.1)
             )
-            xform = rendering.Transform()
-            xform.set_rotation(agent.state.rot[env_index])
-            xform.set_translation(*agent.state.pos[env_index])
-            line.add_attr(xform)
-            line.set_color(*color)
-            geoms.append(line)
 
         return geoms
 

vmas/scenarios/debug/drone.py

@@ -11,7 +11,7 @@
 from vmas.simulator.core import Agent, World
 from vmas.simulator.dynamics.drone import Drone
 from vmas.simulator.scenario import BaseScenario
-from vmas.simulator.utils import Color, ScenarioUtils
+from vmas.simulator.utils import ScenarioUtils
 
 if typing.TYPE_CHECKING:
     from vmas.simulator.rendering import Geom
@@ -93,24 +93,14 @@ def done(self):
         )
 
     def extra_render(self, env_index: int = 0) -> "List[Geom]":
-        from vmas.simulator import rendering
 
         geoms: List[Geom] = []
 
         # Agent rotation
         for agent in self.world.agents:
-            color = Color.BLACK.value
-            line = rendering.Line(
-                (0, 0),
-                (0.1, 0),
-                width=1,
+            geoms.append(
+                ScenarioUtils.plot_entity_rotation(agent, env_index, length=0.1)
             )
-            xform = rendering.Transform()
-            xform.set_rotation(agent.state.rot[env_index])
-            xform.set_translation(*agent.state.pos[env_index])
-            line.add_attr(xform)
-            line.set_color(*color)
-            geoms.append(line)
 
         return geoms
 

vmas/scenarios/debug/kinematic_bicycle.py

@@ -12,7 +12,7 @@
 from vmas.simulator.dynamics.holonomic_with_rot import HolonomicWithRotation
 from vmas.simulator.dynamics.kinematic_bicycle import KinematicBicycle
 from vmas.simulator.scenario import BaseScenario
-from vmas.simulator.utils import Color, ScenarioUtils
+from vmas.simulator.utils import ScenarioUtils
 
 if typing.TYPE_CHECKING:
     from vmas.simulator.rendering import Geom
@@ -99,24 +99,14 @@ def observation(self, agent: Agent):
         )
 
     def extra_render(self, env_index: int = 0) -> "List[Geom]":
-        from vmas.simulator import rendering
 
         geoms: List[Geom] = []
 
         # Agent rotation
         for agent in self.world.agents:
-            color = Color.BLACK.value
-            line = rendering.Line(
-                (0, 0),
-                (0.1, 0),
-                width=1,
+            geoms.append(
+                ScenarioUtils.plot_entity_rotation(agent, env_index, length=0.1)
             )
-            xform = rendering.Transform()
-            xform.set_rotation(agent.state.rot[env_index])
-            xform.set_translation(*agent.state.pos[env_index])
-            line.add_attr(xform)
-            line.set_color(*color)
-            geoms.append(line)
 
         return geoms
 

vmas/simulator/utils.py

@@ -1,8 +1,9 @@
-#  Copyright (c) 2022-2024.
+#  Copyright (c) 2022-2025.
 #  ProrokLab (https://www.proroklab.org/)
 #  All rights reserved.
 import importlib
 import os
+import typing
 import warnings
 from abc import ABC, abstractmethod
 from enum import Enum
@@ -12,6 +13,9 @@
 import torch
 from torch import Tensor
 
+if typing.TYPE_CHECKING:
+    from vmas.simulator.rendering import Geom
+
 _has_matplotlib = importlib.util.find_spec("matplotlib") is not None
 
 X = 0
@@ -50,6 +54,10 @@ class Color(Enum):
     WHITE = (0.75, 0.75, 0.75)
     GRAY = (0.25, 0.25, 0.25)
     BLACK = (0.15, 0.15, 0.15)
+    ORANGE = (1.00, 0.50, 0)
+    PINK = (0.97, 0.51, 0.75)
+    PURPLE = (0.60, 0.31, 0.64)
+    YELLOW = (0.87, 0.87, 0)
 
 
 def override(cls):
@@ -319,3 +327,54 @@ def check_kwargs_consumed(dictionary_of_kwargs: Dict, warn: bool = True):
                 )
             else:
                 raise ValueError(message)
+
+    @staticmethod
+    def render_agent_indices(
+        scenario, env_index: int, start_from: int = 0, exclude: List = None
+    ) -> "List[Geom]":
+        from vmas.simulator import rendering
+
+        aspect_r = scenario.viewer_size[X] / scenario.viewer_size[Y]
+        if aspect_r > 1:
+            dimensional_ratio = (aspect_r, 1)
+        else:
+            dimensional_ratio = (1, 1 / aspect_r)
+
+        geoms = []
+        for i, entity in enumerate(scenario.world.agents):
+            if exclude is not None and entity in exclude:
+                continue
+            i = i + start_from
+            line = rendering.TextLine(
+                text=str(i),
+                font_size=15,
+                x=(
+                    (entity.state.pos[env_index, X] * scenario.viewer_size[X])
+                    / (scenario.viewer_zoom**2 * dimensional_ratio[X] * 2)
+                    + scenario.viewer_size[X] / 2
+                ),
+                y=(
+                    (entity.state.pos[env_index, Y] * scenario.viewer_size[Y])
+                    / (scenario.viewer_zoom**2 * dimensional_ratio[Y] * 2)
+                    + scenario.viewer_size[Y] / 2
+                ),
+            )
+            geoms.append(line)
+        return geoms
+
+    @staticmethod
+    def plot_entity_rotation(entity, env_index: int, length: float = 0.15) -> "Geom":
+        from vmas.simulator import rendering
+
+        color = entity.color
+        line = rendering.Line(
+            (0, 0),
+            (length, 0),
+            width=2,
+        )
+        xform = rendering.Transform()
+        xform.set_rotation(entity.state.rot[env_index])
+        xform.set_translation(*entity.state.pos[env_index])
+        line.add_attr(xform)
+        line.set_color(*color)
+        return line