From d4efe7821cedad8643b0764a172ed6cbb63c5a0a Mon Sep 17 00:00:00 2001 From: dinoSvar Date: Fri, 9 Oct 2026 16:33:02 -0400 Subject: [PATCH 01/10] preserving ADiCo observations and position-query helpers Signed-off-by: dinoSvar --- vmas/scenarios/balance.py | 69 ++++++++++++++++++++++++++++++++++++ vmas/scenarios/navigation.py | 25 ++++++++++--- 2 files changed, 89 insertions(+), 5 deletions(-) diff --git a/vmas/scenarios/balance.py b/vmas/scenarios/balance.py index afb68ab0..4a7efb62 100644 --- a/vmas/scenarios/balance.py +++ b/vmas/scenarios/balance.py @@ -257,6 +257,75 @@ def observation(self, agent: Agent): dim=-1, ) + def observation_from_pos(self, pos: torch.Tensor, env_index: int = None): + """ + Get observation from a given position for balance scenario. + + Args: + pos: Position tensor of shape [batch_size, 2] or [2] + env_index: Index of the environment (optional) + + Returns: + Observation tensor as if an agent were at the given position + """ + # Ensure pos has correct shape + if pos.dim() == 1: + pos = pos.unsqueeze(0) # [2] -> [1, 2] + + batch_size = pos.shape[0] + + # Get states for the specified environment + if env_index is None: + package_pos = self.package.state.pos[0].unsqueeze(0) # [1, 2] + package_vel = self.package.state.vel[0].unsqueeze(0) # [1, 2] + line_pos = self.line.state.pos[0].unsqueeze(0) # [1, 2] + line_vel = self.line.state.vel[0].unsqueeze(0) # [1, 2] + line_ang_vel = self.line.state.ang_vel[0].unsqueeze(0) # [1, 1] + line_rot = self.line.state.rot[0].unsqueeze(0) # [1, 1] + goal_pos = self.package.goal.state.pos[0].unsqueeze(0) # [1, 2] + else: + package_pos = self.package.state.pos[env_index].unsqueeze(0) # [1, 2] + package_vel = self.package.state.vel[env_index].unsqueeze(0) # [1, 2] + line_pos = self.line.state.pos[env_index].unsqueeze(0) # [1, 2] + line_vel = self.line.state.vel[env_index].unsqueeze(0) # [1, 2] + line_ang_vel = self.line.state.ang_vel[env_index].unsqueeze(0) # [1, 1] + line_rot = self.line.state.rot[env_index].unsqueeze(0) # [1, 1] + goal_pos = self.package.goal.state.pos[env_index].unsqueeze(0) # [1, 2] + + # Expand to match batch size + package_pos = package_pos.expand(batch_size, -1) + package_vel = package_vel.expand(batch_size, -1) + line_pos = line_pos.expand(batch_size, -1) + line_vel = line_vel.expand(batch_size, -1) + line_ang_vel = line_ang_vel.expand(batch_size, -1) + line_rot = line_rot.expand(batch_size, -1) + goal_pos = goal_pos.expand(batch_size, -1) + + # Create zero velocity for the hypothetical agent at pos + agent_vel = torch.zeros_like(pos) + + # Match the structure of the observation method: + # [agent.state.pos, agent.state.vel, + # agent.state.pos - package.state.pos, + # agent.state.pos - line.state.pos, + # package.state.pos - goal.state.pos, + # package.state.vel, line.state.vel, line.state.ang_vel, + # line.state.rot % pi] + return torch.cat( + [ + pos, # agent position + agent_vel, # agent velocity (zero for static query point) + pos - package_pos, # relative position to package + pos - line_pos, # relative position to line + package_pos - goal_pos, # package to goal + package_vel, # package velocity + line_vel, # line velocity + line_ang_vel, # line angular velocity + line_rot % torch.pi, # line rotation + ], + dim=-1, + ) + def done(self): return self.on_the_ground + self.world.is_overlapping( self.package, self.package.goal diff --git a/vmas/scenarios/navigation.py b/vmas/scenarios/navigation.py index 876d3363..2f6eae51 100644 --- a/vmas/scenarios/navigation.py +++ b/vmas/scenarios/navigation.py @@ -244,6 +244,8 @@ def agent_reward(self, agent: Agent): return agent.pos_rew def observation(self, agent: Agent): + if self.observe_all_goals and not self.collisions: + return self.observation_from_pos(agent.state.pos) goal_poses = [] if self.observe_all_goals: for a in self.world.agents: @@ -251,11 +253,11 @@ def observation(self, agent: Agent): else: goal_poses.append(agent.state.pos - agent.goal.state.pos) return torch.cat( - [ - agent.state.pos, - agent.state.vel, - ] - + goal_poses + # [ + # agent.state.pos, + # agent.state.vel, + # ] + goal_poses + ( [agent.sensors[0]._max_range - agent.sensors[0].measure()] if self.collisions @@ -264,6 +266,19 @@ def observation(self, agent: Agent): dim=-1, ) + def observation_from_pos(self, pos: Tensor, env_index: typing.Optional[int] = None): + assert self.observe_all_goals and not self.collisions + goal_poses = [] + if self.observe_all_goals: + for a in self.world.agents: + if env_index is None: + goal_poses.append(pos - a.goal.state.pos) + else: + goal_poses.append( + pos - a.goal.state.pos[env_index].unsqueeze(0).expand(pos.shape) + ) + return torch.cat(goal_poses, dim=-1) + def done(self): return torch.stack( [ From f9fadaf9d0586470f69cafe2051e7aed979c9235 Mon Sep 17 00:00:00 2001 From: dinoSvar Date: Fri, 9 Oct 2026 16:37:08 -0400 Subject: [PATCH 02/10] Preserve position-query helpers for three VMAS scenarios Signed-off-by: dinoSvar --- vmas/scenarios/ball_passage.py | 55 +++++++++++++++++++++++++++++++ vmas/scenarios/ball_trajectory.py | 44 +++++++++++++++++++++++++ vmas/scenarios/buzz_wire.py | 42 +++++++++++++++++++++++ 3 files changed, 141 insertions(+) diff --git a/vmas/scenarios/ball_passage.py b/vmas/scenarios/ball_passage.py index b52c3074..cf718f01 100644 --- a/vmas/scenarios/ball_passage.py +++ b/vmas/scenarios/ball_passage.py @@ -287,6 +287,61 @@ def observation(self, agent: Agent): dim=-1, ) + def observation_from_pos(self, pos: torch.Tensor, env_index: int = None): + """ + Get observation from a given position for ball passage scenario. + + Args: + pos: Position tensor of shape [batch_size, 2] or [2] + env_index: Index of the environment (optional) + + Returns: + Observation tensor as if an agent were at the given position + """ + # Ensure pos has correct shape [batch_size, 2] + if pos.dim() == 1: + pos = pos.unsqueeze(0) + + batch_size = pos.shape[0] + device = self.world.device + + # Helper to get the correct state slice + def get_state(entity_state_attr): + if env_index is None: + # If no index, we assume we want the first env or the whole batch + # To match pos batch_size, we take the first and expand + return entity_state_attr[0].unsqueeze(0).expand(batch_size, -1) + return entity_state_attr[env_index].unsqueeze(0).expand(batch_size, -1) + + # Extract states + ball_pos = get_state(self.ball.state.pos) + goal_pos = get_state(self.goal.state.pos) + + # Hypothetical agent velocity is zero for a static position query + agent_vel = torch.zeros_like(pos) + + # Get positions of all "open" passages (where collide is False) + # Note: In VMAS, the list of passages is the same across batch dims, + # but their positions vary per env_index + passage_obs = [] + for passage in self.passages: + if not passage.collide: + p_pos = get_state(passage.state.pos) + passage_obs.append(pos - p_pos) + + # Match the structure of Scenario.observation(): + # [pos, vel, pos - goal, pos - ball, *(pos - open_passages)] + return torch.cat( + [ + pos, + agent_vel, + pos - goal_pos, + pos - ball_pos, + *passage_obs, + ], + dim=-1, + ) + def done(self): return ( ( diff --git a/vmas/scenarios/ball_trajectory.py b/vmas/scenarios/ball_trajectory.py index 7647d13d..ba3c6781 100644 --- a/vmas/scenarios/ball_trajectory.py +++ b/vmas/scenarios/ball_trajectory.py @@ -210,6 +210,50 @@ def observation(self, agent: Agent): dim=-1, ) + def observation_from_pos(self, pos: torch.Tensor, env_index: int = None): + """ + Get observation from a given position for ball trajectory scenario. + + Args: + pos: Position tensor of shape [batch_size, 2] or [2] + env_index: Index of the environment (optional) + + Returns: + Observation tensor as if an agent were at the given position + """ + # Ensure pos has correct shape [batch_size, 2] + if pos.dim() == 1: + pos = pos.unsqueeze(0) + + batch_size = pos.shape[0] + + # Helper to get the correct state slice (single env vs whole batch) + def get_state(entity_state_attr): + if env_index is None: + # Use the first environment's state and expand to match batch_size + return entity_state_attr[0].unsqueeze(0).expand(batch_size, -1) + # Use the specific environment's state + return entity_state_attr[env_index].unsqueeze(0).expand(batch_size, -1) + + # Extract ball state + ball_pos = get_state(self.ball.state.pos) + + # Hypothetical agent velocity is zero for a static position query + agent_vel = torch.zeros_like(pos) + + # Match the structure of Scenario.observation(): + # [agent.pos, agent.vel, agent.pos - ball.pos, agent.pos] + # Note: Your original observation duplicates agent.pos at the end. + return torch.cat( + [ + pos, # agent.state.pos + agent_vel, # agent.state.vel + pos - ball_pos, # agent.state.pos - ball.state.pos + pos, # agent.state.pos (the duplicate from your code) + ], + dim=-1, + ) + def info(self, agent: Agent) -> Dict[str, Tensor]: return { "pos_rew": self.pos_rew, diff --git a/vmas/scenarios/buzz_wire.py b/vmas/scenarios/buzz_wire.py index d353832c..7fc859f8 100644 --- a/vmas/scenarios/buzz_wire.py +++ b/vmas/scenarios/buzz_wire.py @@ -245,6 +245,48 @@ def observation(self, agent: Agent): dim=-1, ) + def observation_from_pos(self, pos: torch.Tensor, env_index: int = None): + """ + Get observation from a given position for the buzz wire scenario. + + Args: + pos: Position tensor of shape [batch_size, 2] or [2] + env_index: Index of the environment (optional) + + Returns: + Observation tensor as if an agent were at the given position + """ + # Ensure pos has correct shape [batch_size, 2] + if pos.dim() == 1: + pos = pos.unsqueeze(0) + + batch_size = pos.shape[0] + + # Helper to get the correct state slice (single env vs whole batch) + def get_state(entity_state_attr): + if env_index is None: + # Use the first environment's state and expand to match batch_size + return entity_state_attr[0].unsqueeze(0).expand(batch_size, -1) + # Use the specific environment's state + return entity_state_attr[env_index].unsqueeze(0).expand(batch_size, -1) + + # Extract goal state + goal_pos = get_state(self.goal.state.pos) + + # Hypothetical agent velocity is zero for a static position query + agent_vel = torch.zeros_like(pos) + + # Match the structure of Scenario.observation(): + # [agent.pos, agent.vel, agent.pos - goal.pos] + return torch.cat( + [ + pos, # agent.state.pos + agent_vel, # agent.state.vel + pos - goal_pos, # agent.state.pos - goal.state.pos + ], + dim=-1, + ) + def done(self): return ( torch.linalg.vector_norm(self.ball.state.pos - self.goal.state.pos, dim=1) From 8d14848df7a1ff722b5e057e4b550efc85907d67 Mon Sep 17 00:00:00 2001 From: dinoSvar Date: Fri, 9 Oct 2026 16:41:43 -0400 Subject: [PATCH 03/10] Preserve Dispersion observations and Reverse Transport position queries --- vmas/scenarios/dispersion.py | 55 ++++++++++++++++++++++++++++- vmas/scenarios/reverse_transport.py | 48 +++++++++++++++++++++++++ 2 files changed, 102 insertions(+), 1 deletion(-) diff --git a/vmas/scenarios/dispersion.py b/vmas/scenarios/dispersion.py index 9229a05f..44075085 100644 --- a/vmas/scenarios/dispersion.py +++ b/vmas/scenarios/dispersion.py @@ -147,10 +147,63 @@ def observation(self, agent: Agent): ) ) return torch.cat( - [agent.state.pos, agent.state.vel, *obs], + obs, dim=-1, ) + def observation_from_pos(self, pos: torch.Tensor, env_index: int = None): + """ + Get observation from a given position. + + Args: + pos: Position tensor of shape [batch_size, 2] or [2] + env_index: Index of the environment (optional) + + Returns: + Observation tensor as if an agent were at the given position + """ + # Ensure pos has correct shape + if pos.dim() == 1: + pos = pos.unsqueeze(0) # [2] -> [1, 2] + + batch_size = pos.shape[0] + + obs = [] + for landmark in self.world.landmarks: + if env_index is None: + # Use all environments + landmark_pos = landmark.state.pos # [batch_dim, 2] + landmark_eaten = landmark.eaten # [batch_dim] + + # If pos batch doesn't match world batch_dim, expand landmark data + if batch_size != landmark_pos.shape[0]: + # Assume we want first env if sizes don't match + landmark_pos = landmark_pos[0:1].expand(batch_size, -1) + landmark_eaten = landmark_eaten[0:1].expand(batch_size) + else: + # Use specific environment + landmark_pos = landmark.state.pos[env_index].unsqueeze(0) # [1, 2] + landmark_eaten = landmark.eaten[env_index].unsqueeze(0) # [1] + + # Expand to match batch size + landmark_pos = landmark_pos.expand(batch_size, -1) # [batch_size, 2] + landmark_eaten = landmark_eaten.expand(batch_size) # [batch_size] + + # Compute relative position from given pos + relative_pos = landmark_pos - pos # [batch_size, 2] + + obs.append( + torch.cat( + [ + relative_pos, + landmark_eaten.to(torch.int).unsqueeze(-1), + ], + dim=-1, + ) + ) + + return torch.cat(obs, dim=-1) + def done(self): return torch.all( torch.stack( diff --git a/vmas/scenarios/reverse_transport.py b/vmas/scenarios/reverse_transport.py index 7dffa9e3..46c73ab6 100644 --- a/vmas/scenarios/reverse_transport.py +++ b/vmas/scenarios/reverse_transport.py @@ -202,6 +202,54 @@ def observation(self, agent: Agent): ], dim=-1, ) + def observation_from_pos(self, pos: torch.Tensor, env_index: int = None): + """ + Get observation from a given position for reverse_transport scenario. + + Args: + pos: Position tensor of shape [batch_size, 2] or [2] + env_index: Index of the environment (optional) + + Returns: + Observation tensor as if an agent were at the given position + """ + # Ensure pos has correct shape + if pos.dim() == 1: + pos = pos.unsqueeze(0) # [2] -> [1, 2] + + batch_size = pos.shape[0] + + # Get package state for the specified environment + if env_index is None: + package_pos = self.package.state.pos[0].unsqueeze(0) # [1, 2] + package_vel = self.package.state.vel[0].unsqueeze(0) # [1, 2] + goal_pos = self.package.goal.state.pos[0].unsqueeze(0) # [1, 2] + else: + package_pos = self.package.state.pos[env_index].unsqueeze(0) # [1, 2] + package_vel = self.package.state.vel[env_index].unsqueeze(0) # [1, 2] + goal_pos = self.package.goal.state.pos[env_index].unsqueeze(0) # [1, 2] + + # Expand to match batch size + package_pos = package_pos.expand(batch_size, -1) + package_vel = package_vel.expand(batch_size, -1) + goal_pos = goal_pos.expand(batch_size, -1) + + # Create zero velocity for the hypothetical agent at pos + agent_vel = torch.zeros_like(pos) + + # Match the structure of the observation method: + # [agent.state.pos, agent.state.vel, package.state.vel, + # package.state.pos - agent.state.pos, package.state.pos - goal.state.pos] + return torch.cat( + [ + pos, # agent position + agent_vel, # agent velocity (zero for static query point) + package_vel, # package velocity + package_pos - pos, # relative position to package + package_pos - goal_pos, # package to goal + ], + dim=-1, + ) def done(self): return self.package.on_goal From ed9cc57fbeb0d9334114dda01e1e1be3744b584f Mon Sep 17 00:00:00 2001 From: dinoSvar Date: Fri, 9 Oct 2026 16:45:06 -0400 Subject: [PATCH 04/10] Preserve agent-relative observations and actions in Multi Give Way --- vmas/scenarios/multi_give_way.py | 91 ++++++++++++++++++++------------ 1 file changed, 56 insertions(+), 35 deletions(-) diff --git a/vmas/scenarios/multi_give_way.py b/vmas/scenarios/multi_give_way.py index 1bd5f485..03ec5444 100644 --- a/vmas/scenarios/multi_give_way.py +++ b/vmas/scenarios/multi_give_way.py @@ -64,6 +64,21 @@ def make_world(self, batch_dim: int, device: torch.device, **kwargs): self.min_collision_distance = 0.005 self.colors = [Color.GREEN, Color.BLUE, Color.RED, Color.GRAY] + self.indices_to_rot_angles = { + 0: torch.pi / 2, + 1: torch.pi, + 2: -torch.pi / 2, + 3: 0, + } + self.indices_to_rot_angles = { + key: torch.full( + (batch_dim,), + value, + device=device, + dtype=torch.float, + ) + for key, value in self.indices_to_rot_angles.items() + } # Add agents for i in range(self.n_agents): @@ -71,11 +86,9 @@ def make_world(self, batch_dim: int, device: torch.device, **kwargs): name=f"agent_{i}", rotatable=False, linear_friction=self.linear_friction, - shape=( - Sphere(radius=self.agent_radius) - if not self.box_agents - else Box(length=self.agent_box_length, width=self.agent_box_width) - ), + shape=Sphere(radius=self.agent_radius) + if not self.box_agents + else Box(length=self.agent_box_length, width=self.agent_box_width), u_range=self.u_range, f_range=self.f_range, render_action=True, @@ -103,6 +116,14 @@ def make_world(self, batch_dim: int, device: torch.device, **kwargs): return world + def rotate_vector(self, agent_index, vector, action: bool): + return TorchUtils.rotate_vector( + vector, + self.indices_to_rot_angles[agent_index] + if not action + else -self.indices_to_rot_angles[agent_index], + ) + def reset_world_at(self, env_index: int = None): for i, agent in enumerate(self.world.agents): agent.controller.reset(env_index) @@ -184,6 +205,9 @@ def reset_world_at(self, env_index: int = None): self.reached_goal[env_index] = False def process_action(self, agent: Agent): + agent.action.u = self.rotate_vector( + self.world.agents.index(agent), agent.action.u, action=True + ) # Clamp square to circle agent.action.u = TorchUtils.clamp_with_norm(agent.action.u, self.u_range) @@ -225,8 +249,7 @@ def reward(self, agent: Agent): self.pos_rew += a.pos_rew self.all_goal_reached = torch.all( - torch.stack([a.on_goal for a in self.world.agents], dim=-1), - dim=-1, + torch.stack([a.on_goal for a in self.world.agents], dim=-1), dim=-1 ) self.final_rew[self.all_goal_reached] = self.final_reward @@ -267,27 +290,29 @@ def reward(self, agent: Agent): ) def observation(self, agent: Agent): - observations = [ - agent.state.pos, - agent.state.vel, + rel_goal = self.rotate_vector( + self.world.agents.index(agent), agent.state.pos - agent.goal.state.pos, - torch.linalg.vector_norm( - agent.state.pos - agent.goal.state.pos, - dim=-1, - ).unsqueeze(-1), + action=False, + ) + + other_agents = [ + self.world.agents[(self.world.agents.index(agent) + i) % self.n_agents] + for i in (1, 2, 3) + ] + rel_poses_to_others = [ + self.rotate_vector( + self.world.agents.index(agent), + agent.state.pos - a.state.pos, + action=False, + ) + for a in other_agents ] - if self.obs_noise > 0: - for i, obs in enumerate(observations): - noise = torch.zeros(*obs.shape, device=self.world.device,).uniform_( - -self.obs_noise, - self.obs_noise, - ) - observations[i] = obs + noise - return torch.cat( - observations, - dim=-1, - ) + return { + "rel_goal": rel_goal, + "rel_agents": torch.stack(rel_poses_to_others, dim=-2), + } def info(self, agent: Agent): return { @@ -361,11 +386,9 @@ def reset_map(self, env_index): landmark.set_pos( torch.tensor( [ - ( - -self.scenario_length / 2 - if i % 2 == 0 - else self.scenario_length / 2 - ), + -self.scenario_length / 2 + if i % 2 == 0 + else self.scenario_length / 2, 0.0, ], dtype=torch.float32, @@ -386,11 +409,9 @@ def reset_map(self, env_index): torch.tensor( [ 0.0, - ( - -self.scenario_length / 2 - if i % 2 == 0 - else self.scenario_length / 2 - ), + -self.scenario_length / 2 + if i % 2 == 0 + else self.scenario_length / 2, ], dtype=torch.float32, device=self.world.device, From 2bc962bd65170a9b63888f44ef5841f7b19b12db Mon Sep 17 00:00:00 2001 From: dinoSvar Date: Fri, 9 Oct 2026 16:56:33 -0400 Subject: [PATCH 05/10] Restore configurable starting angle in Wind Flocking --- vmas/scenarios/wind_flocking.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/vmas/scenarios/wind_flocking.py b/vmas/scenarios/wind_flocking.py index 43737239..9225c320 100644 --- a/vmas/scenarios/wind_flocking.py +++ b/vmas/scenarios/wind_flocking.py @@ -84,6 +84,7 @@ def make_world(self, batch_dim: int, device: torch.device, **kwargs): # Other self.cover_angle_tolerance = kwargs.pop("cover_angle_tolerance", 1) self.horizon = kwargs.pop("horizon", 200) + self.start_angle_range = kwargs.pop("start_angle_range", torch.pi / 8) ScenarioUtils.check_kwargs_consumed(kwargs) self.desired_distance = 1 @@ -154,8 +155,8 @@ def reset_world_at(self, env_index: int = None): device=self.world.device, dtype=torch.float32, ).uniform_( - -torch.pi / 8, - torch.pi / 8, + -self.start_angle_range, + self.start_angle_range, ) start_delta_x = (self.desired_distance / 2) * torch.cos(start_angle) From 622e959cdf49e17fe96ec68e695cb41d348f6c5c Mon Sep 17 00:00:00 2001 From: dinoSvar Date: Fri, 9 Oct 2026 17:05:02 -0400 Subject: [PATCH 06/10] Restore Sampling position queries and fix query coordinates and normalization --- vmas/scenarios/sampling.py | 43 ++++++++++++++++++++++++-------------- 1 file changed, 27 insertions(+), 16 deletions(-) diff --git a/vmas/scenarios/sampling.py b/vmas/scenarios/sampling.py index 86438de8..0c7124d5 100644 --- a/vmas/scenarios/sampling.py +++ b/vmas/scenarios/sampling.py @@ -279,12 +279,20 @@ def reward(self, agent: Agent) -> Tensor: return self.sampling_rew if self.shared_rew else agent.sample def observation(self, agent: Agent) -> Tensor: - observations = [ - agent.state.pos, - agent.state.vel, - agent.sensors[0].measure(), - ] + observations = self.observation_from_pos(agent.state.pos) + + if self.collisions: + observations = torch.cat( + [observations, agents.sensors[0].measure()], dim=-1 + ) + + return observations + def observation_from_pos(self, pos: Tensor, env_index: int = None): + if pos.dim() == 1: + pos = pos.unsqueeze(0) + + samples = [] for delta in [ [self.grid_spacing, 0], [-self.grid_spacing, 0], @@ -295,21 +303,24 @@ def observation(self, agent: Agent) -> Tensor: [-self.grid_spacing, self.grid_spacing], [self.grid_spacing, self.grid_spacing], ]: - pos = agent.state.pos + torch.tensor( + query_pos = pos + torch.tensor( delta, device=self.world.device, dtype=torch.float32, ) - sample = self.sample( - pos, - update_sampled_flag=False, - ).unsqueeze(-1) - observations.append(sample) - - return torch.cat( - observations, - dim=-1, - ) + + if env_index is not None: + sample = self.sample_single_env( + query_pos, env_index=env_index, norm=self.norm + ) + else: + sample = self.sample( + query_pos, update_sampled_flag=False, norm=self.norm + ) + + samples.append(sample.unsqueeze(-1)) + + return torch.cat([pos] + samples, dim=-1) def info(self, agent: Agent) -> Dict[str, Tensor]: return {"agent_sample": agent.sample} From 7145982fe95159f2a84e613b361cdecc908174f9 Mon Sep 17 00:00:00 2001 From: dinoSvar Date: Fri, 9 Oct 2026 17:17:41 -0400 Subject: [PATCH 07/10] Add Discovery position queries using lidar ray geometry --- vmas/scenarios/discovery.py | 50 +++++++++++++++++++++++++++++++++++++ vmas/simulator/core.py | 2 +- 2 files changed, 51 insertions(+), 1 deletion(-) diff --git a/vmas/scenarios/discovery.py b/vmas/scenarios/discovery.py index 2fcc4749..daa7ef77 100644 --- a/vmas/scenarios/discovery.py +++ b/vmas/scenarios/discovery.py @@ -250,6 +250,56 @@ def observation(self, agent: Agent): dim=-1, ) + def observation_from_pos( self, pos: Tensor, env_index: int = None, agent_index: int = 0 ): + env_index = 0 if env_index is None else env_index + agent = self.world.agents[agent_index] + pos = pos.to( + device=self.world.device, dtype=agent.state.pos.dtype + ).reshape(-1, 2) + n_points = pos.shape[0] + + observations = [pos, torch.zeros_like(pos)] + + for sensor in agent.sensors: + entities = [ + entity + for entity in self.world.entities + if entity is not agent and sensor.entity_filter(entity) + ] + + if not entities: + measurements = pos.new_full( + (n_points, sensor._angles.shape[-1]), + sensor._max_range, + ) + else: + sphere_pos = torch.stack( + [entity.state.pos[env_index] for entity in entities] + ).unsqueeze(0).expand(n_points, -1, -1) + + sphere_radius = pos.new_tensor( + [entity.shape.radius for entity in entities] + ).unsqueeze(0).expand(n_points, -1) + + angles = ( + sensor._angles[env_index] + agent.state.rot[env_index] + ).unsqueeze(0).expand(n_points, -1) + + distances = self.world._cast_rays_to_sphere( + sphere_pos, + sphere_radius, + pos, + angles, + sensor._max_range, + ) + measurements = distances.min(dim=-2).values.clamp( + max=sensor._max_range + ) + + observations.append(measurements) + + return torch.cat(observations, dim=-1) + def info(self, agent: Agent) -> Dict[str, Tensor]: info = { "covering_reward": ( diff --git a/vmas/simulator/core.py b/vmas/simulator/core.py index 39af804b..41ffb478 100644 --- a/vmas/simulator/core.py +++ b/vmas/simulator/core.py @@ -1420,7 +1420,7 @@ def _cast_rays_to_sphere( max_range: float, ): batch_size = ray_origin.shape[:-1] - assert batch_size[0] == self.batch_dim + # assert batch_size[0] == self.batch_dim assert ray_origin.shape[-1] == 2 # ray_origin is [*batch_size, 2] assert ( ray_direction.shape[:-1] == batch_size From 030348a9207a654942d44577571a91256f86ff56 Mon Sep 17 00:00:00 2001 From: dinoSvar Date: Fri, 9 Oct 2026 17:22:58 -0400 Subject: [PATCH 08/10] Add Football position queries using the existing observation method --- vmas/scenarios/football.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/vmas/scenarios/football.py b/vmas/scenarios/football.py index cb085e50..4e79275e 100644 --- a/vmas/scenarios/football.py +++ b/vmas/scenarios/football.py @@ -1301,6 +1301,23 @@ def observation( ) return obs + def observation_from_pos( + self, pos: Tensor, env_index: int = None, agent_index: int = 0 + ): + env_index = 0 if env_index is None else env_index + agent = self.world.agents[agent_index] + pos = pos.to( + device=self.world.device, dtype=agent.state.pos.dtype + ).reshape(-1, 2) + + return self.observation( + agent, + agent_pos=pos, + agent_vel=torch.zeros_like(pos), + agent_force=torch.zeros_like(pos), + env_index=env_index, + ) + def observation_base( self, agent_pos, From 9d81b721b00936e40a76ec3601197d436e22387c Mon Sep 17 00:00:00 2001 From: dinoSvar Date: Fri, 9 Oct 2026 17:25:18 -0400 Subject: [PATCH 09/10] Preserve custom Pac Men and Plinko scenarios --- vmas/scenarios/pac_men.py | 678 ++++++++++++++++++++++++++++++++++++++ vmas/scenarios/plinko.py | 447 +++++++++++++++++++++++++ 2 files changed, 1125 insertions(+) create mode 100644 vmas/scenarios/pac_men.py create mode 100644 vmas/scenarios/plinko.py diff --git a/vmas/scenarios/pac_men.py b/vmas/scenarios/pac_men.py new file mode 100644 index 00000000..779045f6 --- /dev/null +++ b/vmas/scenarios/pac_men.py @@ -0,0 +1,678 @@ +# Copyright (c) 2022-2024. +# ProrokLab (https://www.proroklab.org/) +# All rights reserved. + + +import torch + +from vmas import render_interactively + +from vmas.simulator.core import Agent, Landmark, Line, Sphere, World +from vmas.simulator.scenario import BaseScenario +from vmas.simulator.utils import Color, ScenarioUtils, TorchUtils, X, Y + + +class Scenario(BaseScenario): + def make_world(self, batch_dim: int, device: torch.device, **kwargs): + + self.corridors_length = kwargs.get("corridors_length", 1) + self.lrud_ratio = kwargs.get("lrud_ratio", (2, 2, 3, 1)) + self.shared_rew = kwargs.get("shared_rew", True) + self.observe_position = kwargs.get("observe_position", False) + self.spawn_same_pos = kwargs.get("spawn_same_pos", True) + + self.left_corridor_multplier = self.lrud_ratio[0] + self.right_corridor_multplier = self.lrud_ratio[1] + self.up_corridor_multplier = self.lrud_ratio[2] + self.down_corridor_multplier = self.lrud_ratio[3] + + self.viewer_zoom = 3.1 + + self.n_agents = 4 + self.agent_radius = 0.16 + + # Make world + world = World( + batch_dim, + device, + ) + self.spawn_map(world) + + self.xdim = kwargs.get("xdim", 5) + self.ydim = kwargs.get("ydim", 5) + self.grid_spacing = kwargs.get("grid_spacing", self.agent_radius * 2) + + self.plot_grid = True + self.alpha_plot: float = 0.5 + + self.x_semidim = self.room_semidim - self.agent_radius + self.y_semidim = self.room_semidim - self.agent_radius + self.n_x_cells = int((2 * self.room_semidim) / self.grid_spacing) + self.n_y_cells = int((2 * self.room_semidim) / self.grid_spacing) + + self.distributions = [] + for _ in range(4): + dist = torch.distributions.uniform.Uniform( + low=torch.tensor([-self.room_semidim], device=device, dtype=torch.float) + .unsqueeze(0) + .expand(batch_dim, 2) + .clone(), + high=torch.tensor([self.room_semidim], device=device, dtype=torch.float) + .unsqueeze(0) + .expand(batch_dim, 2) + .clone(), + ) + self.distributions.append(dist) + + dist.sampled = torch.zeros( + (batch_dim, self.n_x_cells, self.n_y_cells), + device=device, + dtype=torch.bool, + ) + + self.colors = [Color.GREEN, Color.BLUE, Color.RED, Color.GRAY] + + # Add agents + for i in range(self.n_agents): + agent = Agent( + name=f"agent_{i}", + rotatable=False, + shape=Sphere(radius=self.agent_radius), + render_action=True, + color=self.colors[i], + collide=True, + collision_filter=lambda e: isinstance(e.shape, Line), + ) + agent.sample = torch.zeros(batch_dim, device=device, dtype=torch.float) + world.add_agent(agent) + + return world + + def reset_world_at(self, env_index: int = None): + + ScenarioUtils.spawn_entities_randomly( + self.world.agents, + self.world, + env_index, + min_dist_between_entities=0, + x_bounds=( + -(self.inner_box_semidim - self.agent_radius) + if not self.spawn_same_pos + else 0, + (self.inner_box_semidim - self.agent_radius) + if not self.spawn_same_pos + else 0, + ), + y_bounds=( + -(self.inner_box_semidim - self.agent_radius) + if not self.spawn_same_pos + else 0, + (self.inner_box_semidim - self.agent_radius) + if not self.spawn_same_pos + else 0, + ), + ) + + self.reset_map(env_index) + self.reset_distribution(env_index) + + def reset_distribution(self, env_index): + for dist in self.distributions: + if env_index is None: + dist.sampled = torch.zeros_like(dist.sampled) + else: + dist.sampled = TorchUtils.where_from_index( + env_index, False, dist.sampled + ) + + def translate_pos(self, pos, room_index, inverse=False): + if room_index == 0: + shift = ( + self.left_corridor_length + self.room_semidim + self.inner_box_semidim + ) + pos[..., X] = pos[..., X] - (shift if not inverse else -shift) + if room_index == 1: + shift = ( + self.right_corridor_length + self.room_semidim + self.inner_box_semidim + ) + pos[..., X] = pos[..., X] + (shift if not inverse else -shift) + if room_index == 2: + shift = self.up_corridor_length + self.room_semidim + self.inner_box_semidim + pos[..., Y] = pos[..., Y] + (shift if not inverse else -shift) + if room_index == 3: + shift = ( + self.down_corridor_length + self.room_semidim + self.inner_box_semidim + ) + pos[..., Y] = pos[..., Y] - (shift if not inverse else -shift) + return pos + + def sample( + self, + pos, + update_sampled_flag: bool = False, + ): + values = [] + for i, dist in enumerate(self.distributions): + trans_pos = self.translate_pos(pos.clone(), i, inverse=True) + out_of_bounds = ( + (trans_pos < -self.room_semidim) + (trans_pos > self.room_semidim) + ).any(-1) + if out_of_bounds.all(): + continue + trans_pos[:, X].clamp_(-self.x_semidim, self.x_semidim) + trans_pos[:, Y].clamp_(-self.y_semidim, self.y_semidim) + + index = trans_pos / self.grid_spacing + index[:, X] += self.n_x_cells / 2 + index[:, Y] += self.n_y_cells / 2 + index = index.to(torch.long) + v = dist.log_prob(trans_pos).exp().prod(-1) + v = v * (2 * self.room_semidim) ** 2 # Make it 1 + + sampled = dist.sampled[ + torch.arange(self.world.batch_dim), index[:, X], index[:, Y] + ] + invalid_sample = sampled + out_of_bounds + v[invalid_sample] = 0 + if update_sampled_flag: + dist.sampled[ + torch.arange(self.world.batch_dim), index[:, X], index[:, Y] + ] += ~invalid_sample + values.append(v) + if len(values) == 0: + return torch.zeros_like(pos).sum(dim=-1) + return torch.stack(values, dim=-1).sum(-1) + + def sample_single_env( + self, + pos, + env_index, + ): + pos = pos.view(-1, self.world.dim_p) + values = [] + for i, dist in enumerate(self.distributions): + trans_pos = self.translate_pos(pos.clone(), i, inverse=True) + out_of_bounds = ( + (trans_pos < -self.room_semidim) + (trans_pos > self.room_semidim) + ).any(-1) + + if out_of_bounds.all(): + continue + trans_pos[:, X].clamp_(-self.x_semidim, self.x_semidim) + trans_pos[:, Y].clamp_(-self.y_semidim, self.y_semidim) + + index = trans_pos / self.grid_spacing + index[:, X] += self.n_x_cells / 2 + index[:, Y] += self.n_y_cells / 2 + index = index.to(torch.long) + + trans_pos = trans_pos.unsqueeze(1).expand( + trans_pos.shape[0], self.world.batch_dim, 2 + ) + + v = dist.log_prob(trans_pos).exp().prod(-1)[:, env_index] + v = v * (2 * self.room_semidim) ** 2 # Make it 1 + + sampled = dist.sampled[env_index, index[:, 0], index[:, 1]] + + v[sampled + out_of_bounds] = 0 + + values.append(v) + + if len(values) == 0: + return torch.zeros_like(pos).sum(dim=-1) + return torch.stack(values, dim=-1).sum(-1) + + def reward(self, agent: Agent): + is_first = self.world.agents.index(agent) == 0 + if is_first: + for a in self.world.agents: + a.sample = self.sample(a.state.pos, update_sampled_flag=True) + self.sampling_rew = torch.stack( + [a.sample for a in self.world.agents], dim=-1 + ).sum(-1) + + return self.sampling_rew if self.shared_rew else agent.sample + + def observation(self, agent: Agent): + observations = self.observation_from_pos(agent.state.pos) + return observations + + def observation_from_pos(self, pos, env_index=None): + samples = [] + in_pos = pos + for delta in [ + [self.grid_spacing, 0], + [-self.grid_spacing, 0], + [0, self.grid_spacing], + [0, -self.grid_spacing], + [-self.grid_spacing, -self.grid_spacing], + [self.grid_spacing, -self.grid_spacing], + [-self.grid_spacing, self.grid_spacing], + [self.grid_spacing, self.grid_spacing], + ]: + pos = in_pos + torch.tensor( + delta, + device=self.world.device, + dtype=torch.float32, + ) + if env_index is not None: + samples.append( + self.sample_single_env(pos, env_index=env_index).unsqueeze(-1) + ) + else: + samples.append( + self.sample(pos, update_sampled_flag=False).unsqueeze(-1) + ) + if self.observe_position: + samples.append(pos) + + return torch.cat( + samples, + dim=-1, + ) + + def info(self, agent: Agent): + return {"agent_sample": agent.sample} + + def density_for_plot(self, env_index): + def f(x): + sample = self.sample_single_env( + torch.tensor(x, dtype=torch.float32, device=self.world.device), + env_index=env_index, + ) + return sample + + return f + + def extra_render(self, env_index: int = 0): + + from vmas.simulator.rendering import render_function_util + + geoms = [] + + # Function + for x_bounds, y_bounds in zip(self.x_bounds, self.y_bounds): + res = render_function_util( + f=self.density_for_plot(env_index=env_index), + plot_range=(x_bounds, y_bounds), + cmap_alpha=self.alpha_plot, + precision=self.agent_radius / 2, + cmap_range=(0, 1), + ) + geoms.append(res) + + return geoms + + def spawn_map(self, world: World): + self.corridors_width = self.agent_radius * 4 + self.inner_box_semidim = self.corridors_width / 2 + 1e-5 + + self.room_semidim = self.corridors_width / 2 + 1e-6 + + self.left_corridor_length = self.corridors_length * self.left_corridor_multplier + self.right_corridor_length = ( + self.corridors_length * self.right_corridor_multplier + ) + self.up_corridor_length = self.corridors_length * self.up_corridor_multplier + self.down_corridor_length = self.corridors_length * self.down_corridor_multplier + + self.inner_box_lines_length = self.inner_box_semidim - self.corridors_width / 2 + + self.x_bounds = [ + ( + -( + self.left_corridor_length + + self.inner_box_semidim + + self.room_semidim * 2 + ), + -(self.left_corridor_length + self.inner_box_semidim), + ), + ( + +(self.right_corridor_length + self.inner_box_semidim), + +( + self.right_corridor_length + + self.inner_box_semidim + + self.room_semidim * 2 + ), + ), + ( + -self.room_semidim, + self.room_semidim, + ), + ( + -self.room_semidim, + self.room_semidim, + ), + ] + self.y_bounds = [ + ( + -self.room_semidim, + self.room_semidim, + ), + ( + -self.room_semidim, + self.room_semidim, + ), + ( + (self.up_corridor_length + self.inner_box_semidim), + self.up_corridor_length + + self.inner_box_semidim + + self.room_semidim * 2, + ), + ( + -( + self.down_corridor_length + + self.inner_box_semidim + + self.room_semidim * 2 + ), + -(self.down_corridor_length + self.inner_box_semidim), + ), + ] + + self.inner_box_lines = [] + for i in range(8): + landmark = Landmark( + name=f"inner_wall_{i}", + collide=True, + shape=Line(length=self.inner_box_lines_length), + color=Color.BLACK, + ) + self.inner_box_lines.append(landmark) + world.add_landmark(landmark) + self.corridor_lines = [] + self.corridors_lengths = [ + self.left_corridor_length, + self.right_corridor_length, + self.up_corridor_length, + self.down_corridor_length, + ] + for i in range(8): + landmark = Landmark( + name=f"corridor_{i}", + collide=True, + shape=Line(length=self.corridors_lengths[i % 4]), + color=Color.BLACK, + ) + self.corridor_lines.append(landmark) + world.add_landmark(landmark) + self.room_lines = [] + for i in range(12): + landmark = Landmark( + name=f"room_line_{i}", + collide=True, + shape=Line(length=self.room_semidim * 2), + color=Color.BLACK, + ) + self.room_lines.append(landmark) + world.add_landmark(landmark) + self.closing_lines = [] + for i in range(8): + landmark = Landmark( + name=f"closing_line_{i}", + collide=True, + shape=Line(length=self.room_semidim - self.corridors_width / 2), + color=Color.BLACK, + ) + self.closing_lines.append(landmark) + world.add_landmark(landmark) + + def reset_map(self, env_index): + self.reset_inner_walls(env_index) + self.reset_corridors(env_index) + self.reset_room_lines(env_index) + self.reset_closing_lines(env_index) + + def reset_closing_lines(self, env_index): + for i in [2, 3, 6, 7]: + self.closing_lines[i].set_pos( + torch.tensor( + [ + -(self.room_semidim / 2 + self.corridors_width / 4) + if i in [2, 3] + else (self.room_semidim / 2 + self.corridors_width / 4), + (self.inner_box_semidim + self.corridors_lengths[i % 4]) + * (-1 if i % 2 else 1), + ], + dtype=torch.float32, + device=self.world.device, + ), + batch_index=env_index, + ) + for i in [0, 1, 4, 5]: + self.closing_lines[i].set_pos( + torch.tensor( + [ + (self.inner_box_semidim + self.corridors_lengths[i % 4]) + * (-1 if i % 2 else 1), + -(self.room_semidim / 2 + self.corridors_width / 4) + if i in [0, 1] + else (self.room_semidim / 2 + self.corridors_width / 4), + ], + dtype=torch.float32, + device=self.world.device, + ), + batch_index=env_index, + ) + self.closing_lines[i].set_rot( + torch.tensor( + [torch.pi / 2], + dtype=torch.float32, + device=self.world.device, + ), + batch_index=env_index, + ) + + def reset_room_lines(self, env_index): + for i in [2, 3]: + self.room_lines[i].set_pos( + torch.tensor( + [ + 0, + ( + self.inner_box_semidim + + self.corridors_lengths[i % 4] + + self.room_semidim * 2 + ) + * (-1 if i % 2 else 1), + ], + dtype=torch.float32, + device=self.world.device, + ), + batch_index=env_index, + ) + + for i in [0, 1]: + self.room_lines[i].set_pos( + torch.tensor( + [ + ( + self.inner_box_semidim + + self.corridors_lengths[i % 4] + + self.room_semidim * 2 + ) + * (-1 if i % 2 else 1), + 0, + ], + dtype=torch.float32, + device=self.world.device, + ), + batch_index=env_index, + ) + self.room_lines[i].set_rot( + torch.tensor( + [torch.pi / 2], + dtype=torch.float32, + device=self.world.device, + ), + batch_index=env_index, + ) + for i in [6, 7, 10, 11]: + self.room_lines[i].set_pos( + torch.tensor( + [ + -self.room_semidim if i in [6, 7] else self.room_semidim, + ( + self.inner_box_semidim + + self.corridors_lengths[i % 4] + + self.room_semidim + ) + * (-1 if i % 2 else 1), + ], + dtype=torch.float32, + device=self.world.device, + ), + batch_index=env_index, + ) + self.room_lines[i].set_rot( + torch.tensor( + [torch.pi / 2], + dtype=torch.float32, + device=self.world.device, + ), + batch_index=env_index, + ) + for i in [4, 5, 8, 9]: + self.room_lines[i].set_pos( + torch.tensor( + [ + ( + self.inner_box_semidim + + self.corridors_lengths[i % 4] + + self.room_semidim + ) + * (-1 if i % 2 else 1), + -self.room_semidim if i in [4, 5] else self.room_semidim, + ], + dtype=torch.float32, + device=self.world.device, + ), + batch_index=env_index, + ) + + def reset_inner_walls(self, env_index): + for i, landmark in enumerate(self.inner_box_lines): + if i < 4: + landmark.set_pos( + torch.tensor( + [ + -self.inner_box_semidim + if i in [0, 2] + else self.inner_box_semidim, + self.inner_box_lines_length / 2 + self.corridors_width / 2 + if i in [0, 1] + else -( + self.inner_box_lines_length / 2 + + self.corridors_width / 2 + ), + ], + dtype=torch.float32, + device=self.world.device, + ), + batch_index=env_index, + ) + landmark.set_rot( + torch.tensor( + [torch.pi / 2], + dtype=torch.float32, + device=self.world.device, + ), + batch_index=env_index, + ) + else: + landmark.set_pos( + torch.tensor( + [ + self.inner_box_lines_length / 2 + self.corridors_width / 2 + if i in [4, 6] + else -( + self.inner_box_lines_length / 2 + + self.corridors_width / 2 + ), + -self.inner_box_semidim + if i in [4, 5] + else self.inner_box_semidim, + ], + dtype=torch.float32, + device=self.world.device, + ), + batch_index=env_index, + ) + + def reset_corridors(self, env_index): + for i in [0, 4]: + self.corridor_lines[i].set_pos( + torch.tensor( + [ + -(self.inner_box_semidim + self.corridors_lengths[i % 4] / 2), + -self.corridors_width / 2 + if i == 0 + else self.corridors_width / 2, + ], + dtype=torch.float32, + device=self.world.device, + ), + batch_index=env_index, + ) + for i in [1, 5]: + self.corridor_lines[i].set_pos( + torch.tensor( + [ + (self.inner_box_semidim + self.corridors_lengths[i % 4] / 2), + -self.corridors_width / 2 + if i == 1 + else self.corridors_width / 2, + ], + dtype=torch.float32, + device=self.world.device, + ), + batch_index=env_index, + ) + for i in [2, 6]: + self.corridor_lines[i].set_pos( + torch.tensor( + [ + -self.corridors_width / 2 + if i == 2 + else self.corridors_width / 2, + (self.inner_box_semidim + self.corridors_lengths[i % 4] / 2), + ], + dtype=torch.float32, + device=self.world.device, + ), + batch_index=env_index, + ) + self.corridor_lines[i].set_rot( + torch.tensor( + [torch.pi / 2], + dtype=torch.float32, + device=self.world.device, + ), + batch_index=env_index, + ) + for i in [3, 7]: + self.corridor_lines[i].set_pos( + torch.tensor( + [ + -self.corridors_width / 2 + if i == 3 + else self.corridors_width / 2, + -(self.inner_box_semidim + self.corridors_lengths[i % 4] / 2), + ], + dtype=torch.float32, + device=self.world.device, + ), + batch_index=env_index, + ) + self.corridor_lines[i].set_rot( + torch.tensor( + [torch.pi / 2], + dtype=torch.float32, + device=self.world.device, + ), + batch_index=env_index, + ) + + +if __name__ == "__main__": + render_interactively(__file__, control_two_agents=True) diff --git a/vmas/scenarios/plinko.py b/vmas/scenarios/plinko.py new file mode 100644 index 00000000..3796a1e2 --- /dev/null +++ b/vmas/scenarios/plinko.py @@ -0,0 +1,447 @@ +# Copyright (c) 2023. +# ProrokLab (https://www.proroklab.org/) +# All rights reserved. +import math +import typing + +import torch +from vmas import render_interactively +from vmas.simulator.core import Agent, Landmark, World, Sphere, Line +from vmas.simulator.scenario import BaseScenario +from vmas.simulator.sensors import Lidar +from vmas.simulator.utils import Color, Y, X + +if typing.TYPE_CHECKING: + pass + + +class Scenario(BaseScenario): + def make_world(self, batch_dim: int, device: torch.device, **kwargs): + self.plot_grid = False + self.viewer_zoom = 1.7 + self.agents_start_y = 2 + self.min_collision_distance = 0.005 + + self.n_agents = kwargs.get("n_agents", 2) + self.u_range = kwargs.get("u_range", 1) + self.n_obstacle_rows = kwargs.get("n_obstacle_rows", 5) + self.n_obstacle_columns = kwargs.get("n_obstacle_columns", 7) + self.start_with_sidegaps = kwargs.get("start_with_sidegaps", False) + self.passages_wiggle_room_radius = kwargs.get("passages_wiggle_room_radius", 0) + self.agents_distance_from_obstacle_layer = kwargs.get( + "agents_distance_from_obstacle_layer", 0.7 + ) + self.agent_radius = kwargs.get("agent_radius", 0.1) + self.obstacle_radius = kwargs.get("obstacle_radius", 0.1) + self.agents_x_spawning_distance = kwargs.get("agents_x_spawning_distance", 0.1) + self.rand_agent_order = kwargs.get("rand_agent_order", True) + self.lidar_range = kwargs.get("lidar_range", 0.3) + self.lidar = kwargs.get("lidar", False) + self.boundaries = kwargs.get("boundaries", False) + self.spawn_vertically = kwargs.get("spawn_vertically", True) + self.random_x_start = kwargs.get("random_x_start", False) + self.goal_2d = kwargs.get("goal_2d", False) + self.side_space = kwargs.get("side_space", 4 * self.agent_radius) + self.goal_distance_from_obstacles = kwargs.get( + "goal_distance_from_obstacles", + self.agents_distance_from_obstacle_layer + 0.5, + ) + + self.shared_rew = kwargs.get("shared_rew", True) + self.pos_shaping_factor = kwargs.get("pos_shaping_factor", 1) + self.final_reward = kwargs.get("final_reward", 0) + self.agent_collision_penalty = kwargs.get("agent_collision_penalty", -1) + + self._compute_startup_positions_and_sizes(device) + + # Make world + world = World(batch_dim, device, substeps=3, x_semidim=self.world_x_semidim) + + known_colors = [ + (0.22, 0.49, 0.72), + (1.00, 0.50, 0), + (0.30, 0.69, 0.29), + (0.97, 0.51, 0.75), + (0.60, 0.31, 0.64), + (0.89, 0.10, 0.11), + (0.87, 0.87, 0), + ] + colors = torch.randn( + (max(self.n_agents - len(known_colors), 0), 3), device=device + ) + entity_filter_agents = lambda e: isinstance(e, Landmark) + + # Add agents + for i in range(self.n_agents): + color = ( + known_colors[i] + if i < len(known_colors) + else colors[i - len(known_colors)] + ) + + # Constraint: all agents have same action range and multiplier + agent = Agent( + name=f"agent_{i}", + collide=True, + color=color, + shape=Sphere(radius=self.agent_radius), + render_action=True, + u_range=self.u_range, + sensors=[ + Lidar( + world, + n_rays=8, + max_range=self.lidar_range, + entity_filter=entity_filter_agents, + ), + ], + ) + agent.pos_rew = torch.zeros(batch_dim, device=device) + agent.agent_collision_rew = agent.pos_rew.clone() + world.add_agent(agent) + + # Add obstacles + for i in range(self.n_obstacles): + obstacle = Landmark( + name=f"obstacle_{i}", + collide=True, + movable=False, + rotatable=False, + shape=Sphere(radius=self.obstacle_radius), + ) + world.add_landmark(obstacle) + + # Lines + if self.boundaries: + self.lines = [] + for i in range(2): + line = Landmark( + name=f"line_{i}", + collide=True, + movable=False, + rotatable=False, + shape=Line(length=2 * self.flock_semidist + 2 * self.agent_radius), + ) + self.lines.append(line) + world.add_landmark(line) + line = Landmark( + name=f"line_{2}", + collide=True, + movable=False, + rotatable=False, + shape=Line(length=2 * self.agent_radius), + ) + self.lines.append(line) + world.add_landmark(line) + + self.pos_rew = torch.zeros(batch_dim, device=device) + self.final_rew = self.pos_rew.clone() + + return world + + def _reset_lines(self, env_index): + for i, line in enumerate(self.lines[:-1]): + line.set_pos( + torch.tensor( + [ + -self.agent_radius if i == 0 else self.agent_radius, + self.agents_start_y + self.flock_semidist, + ], + device=self.world.device, + dtype=torch.float, + ), + batch_index=env_index, + ) + line.set_rot( + torch.tensor( + [torch.pi / 2], + device=self.world.device, + dtype=torch.float, + ), + batch_index=env_index, + ) + self.lines[-1].set_pos( + torch.tensor( + [ + 0.0, + self.agents_start_y + 2 * self.flock_semidist + self.agent_radius, + ], + device=self.world.device, + dtype=torch.float, + ), + batch_index=env_index, + ) + + def _compute_startup_positions_and_sizes(self, device): + # How many obstacles + self.n_obstacles = (self.n_obstacle_rows // 2) * ( + self.n_obstacle_columns + 1 + ) + (self.n_obstacle_rows // 2) * self.n_obstacle_columns + if self.n_obstacle_rows % 2: + self.n_obstacles += ( + self.n_obstacle_columns + if self.start_with_sidegaps + else self.n_obstacle_columns + 1 + ) + # Agents spacing + self.agents_x_spawning_distance += self.agent_radius * 2 + # Gap radius + self.gap_radius = self.passages_wiggle_room_radius + self.agent_radius + # Side barriers radius + self.world_x_semidim = ( + self.obstacle_radius + self.gap_radius + ) * self.n_obstacle_columns + self.side_space + # X spacing of obstacles + self.obstacle_x_spacing = self.obstacle_radius * 2 + self.gap_radius * 2 + # Y spacing of obstacles + self.y_distance_rows = math.sqrt( + (2 * self.obstacle_radius + 2 * self.gap_radius) ** 2 + - (self.gap_radius + self.obstacle_radius) ** 2 + ) + self.y_obstacles_finish = ( + self.agents_start_y + - self.agents_distance_from_obstacle_layer + - (self.n_obstacle_rows - 1) * self.y_distance_rows + ) + + # Goal line + self.goal = torch.tensor( + [ + 0, + self.y_obstacles_finish - self.goal_distance_from_obstacles - 0.4, + ], + device=device, + dtype=torch.float, + ) + # Flock semidist + self.flock_semidist = ( + ( + self.agents_x_spawning_distance / 2 + + self.agents_x_spawning_distance * ((self.n_agents // 2) - 1) + ) + if not self.n_agents % 2 + else self.agents_x_spawning_distance * (self.n_agents // 2) + ) + + def reset_world_at(self, env_index: int = None): + if self.boundaries: + self._reset_lines(env_index) + start_pos = torch.cat( + [ + torch.zeros( + (1, 1) if env_index is not None else (self.world.batch_dim, 1), + device=self.world.device, + dtype=torch.float32, + ).uniform_( + ( + -self.world_x_semidim / 2 + + (self.flock_semidist if not self.spawn_vertically else 0) + if self.random_x_start + else 0 + ), + ( + ( + self.world_x_semidim / 2 - self.flock_semidist + if not self.spawn_vertically + else 0 + ) + if self.random_x_start + else 0 + ), + ), + torch.zeros( + (1, 1) if env_index is not None else (self.world.batch_dim, 1), + device=self.world.device, + dtype=torch.float32, + ).uniform_( + self.agents_start_y, + self.agents_start_y, + ), + ], + dim=-1, + ) + + order = torch.randperm(self.n_agents).tolist() + agents = [self.world.agents[i] for i in order] + for i, agent in enumerate( + agents if self.rand_agent_order else self.world.agents + ): + agent_pos = start_pos.clone() + if self.spawn_vertically: + agent_pos[..., Y] += self.agents_x_spawning_distance * i + else: + agent_pos[..., X] += ( + -self.flock_semidist + self.agents_x_spawning_distance * i + ) + + agent.set_pos( + agent_pos, + batch_index=env_index, + ) + if env_index is None: + agent.pos_shaping = ( + torch.abs(agent.state.pos[..., Y] - self.goal[Y]) + if not self.goal_2d + else torch.linalg.vector_norm(agent.state.pos - self.goal, dim=-1) + * self.pos_shaping_factor + ) + + else: + agent.pos_shaping[env_index] = ( + torch.abs(agent.state.pos[env_index, Y] - self.goal[Y]) + if not self.goal_2d + else torch.linalg.vector_norm( + agent.state.pos[env_index] - self.goal, dim=-1 + ) + * self.pos_shaping_factor + ) + + obstacle_total_index = 0 + for row in range(self.n_obstacle_rows): + if not (row % 2): + n_obstacles_this_row = ( + self.n_obstacle_columns + if self.start_with_sidegaps + else self.n_obstacle_columns + 1 + ) + else: + n_obstacles_this_row = ( + self.n_obstacle_columns + 1 + if self.start_with_sidegaps + else self.n_obstacle_columns + ) + for obstacle_index in range(n_obstacles_this_row): + if not n_obstacles_this_row % 2: + pos = [ + -self.obstacle_x_spacing / 2 + - self.obstacle_x_spacing * ((n_obstacles_this_row // 2) - 1) + + self.obstacle_x_spacing * obstacle_index, + self.agents_start_y + - self.agents_distance_from_obstacle_layer + - self.y_distance_rows * row, + ] + else: + pos = [ + -self.obstacle_x_spacing * (n_obstacles_this_row // 2) + + self.obstacle_x_spacing * obstacle_index, + self.agents_start_y + - self.agents_distance_from_obstacle_layer + - self.y_distance_rows * row, + ] + obstacle = self.world.landmarks[obstacle_total_index] + obstacle.set_pos( + torch.tensor( + pos, + device=self.world.device, + ), + batch_index=env_index, + ) + obstacle_total_index += 1 + assert obstacle_total_index == self.n_obstacles + + def reward(self, agent: Agent): + is_first = agent == self.world.agents[0] + + if is_first: + self.pos_rew[:] = 0 + self.final_rew[:] = 0 + + for a in self.world.agents: + self.pos_rew += self.agent_reward(a) + a.agent_collision_rew[:] = 0 + + self.all_goal_reached = torch.all( + torch.stack([a.on_goal for a in self.world.agents], dim=-1), dim=-1 + ) + + self.final_rew[self.all_goal_reached] = self.final_reward + + for i, a in enumerate(self.world.agents): + for j, b in enumerate(self.world.agents): + if i <= j: + continue + if self.world.collides(a, b) and self.agent_collision_penalty < 0: + safe_pos = a.state.pos[..., Y] < self.y_obstacles_finish + + distance = self.world.get_distance(a, b) + close = distance <= self.min_collision_distance + mask = close * ~safe_pos + + a.agent_collision_rew[mask] += self.agent_collision_penalty + b.agent_collision_rew[mask] += self.agent_collision_penalty + + pos_reward = self.pos_rew if self.shared_rew else agent.pos_rew + return pos_reward + self.final_rew + agent.agent_collision_rew + + def agent_reward(self, agent: Agent): + agent.distance_to_goal = ( + torch.abs(agent.state.pos[..., Y] - self.goal[Y]) + if not self.goal_2d + else torch.linalg.vector_norm(agent.state.pos - self.goal, dim=-1) + ) + + agent.on_goal = agent.distance_to_goal < agent.shape.radius + + pos_shaping = agent.distance_to_goal * self.pos_shaping_factor + agent.pos_rew = agent.pos_shaping - pos_shaping + agent.pos_shaping = pos_shaping + return agent.pos_rew + + def observation(self, agent: Agent): + return torch.cat( + [ + agent.state.pos[..., Y].unsqueeze(-1), + ] + + ([self.lidar_range - agent.sensors[0].measure()] if self.lidar else []), + dim=-1, + ) + + def extra_render(self, env_index: int = 0): + from vmas.simulator import rendering + + geoms = [] + + # Perimeter + for i in range(2): + geom = Line(length=50).get_geometry() + xform = rendering.Transform() + geom.add_attr(xform) + + dist = ( + self.world_x_semidim + + self.agent_radius + + self.passages_wiggle_room_radius + ) + + xform.set_translation( + -dist if i == 0 else dist, + 0.0, + ) + xform.set_rotation(torch.pi / 2) + color = Color.BLACK.value + if isinstance(color, torch.Tensor) and len(color.shape) > 1: + color = color[env_index] + geom.set_color(*color) + geoms.append(geom) + + # Y goal + + geom = Line(length=dist * 2).get_geometry() + xform = rendering.Transform() + geom.add_attr(xform) + + xform.set_translation(*self.goal.tolist()) + color = Color.RED.value + if isinstance(color, torch.Tensor) and len(color.shape) > 1: + color = color[env_index] + geom.set_color(*color) + geoms.append(geom) + + return geoms + + +if __name__ == "__main__": + render_interactively( + __file__, + control_two_agents=True, + ) From 06509b93b62af7cb76c1c33ef68735ae8a419cfa Mon Sep 17 00:00:00 2001 From: dinoSvar Date: Fri, 9 Oct 2026 17:26:17 -0400 Subject: [PATCH 10/10] Preserve custom empty debug scenario --- vmas/scenarios/debug/empty.py | 65 +++++++++++++++++++++++++++++++++++ 1 file changed, 65 insertions(+) create mode 100644 vmas/scenarios/debug/empty.py diff --git a/vmas/scenarios/debug/empty.py b/vmas/scenarios/debug/empty.py new file mode 100644 index 00000000..ee667053 --- /dev/null +++ b/vmas/scenarios/debug/empty.py @@ -0,0 +1,65 @@ + +import torch + +from vmas import render_interactively + +from vmas.simulator.core import Agent, Sphere, World +from vmas.simulator.scenario import BaseScenario +from vmas.simulator.utils import Color, ScenarioUtils + + +class Scenario(BaseScenario): + def make_world(self, batch_dim: int, device: torch.device, **kwargs): + + self.n_agents = 4 + self.agent_radius = 0.16 + + # Make world + world = World(batch_dim, device, x_semidim=0, y_semidim=0) + + self.colors = [Color.GREEN, Color.BLUE, Color.RED, Color.GRAY] + + # Add agents + for i in range(self.n_agents): + agent = Agent( + name=f"agent_{i}", + rotatable=False, + shape=Sphere(radius=self.agent_radius), + render_action=True, + color=self.colors[i], + collide=False, + ) + world.add_agent(agent) + + return world + + def reset_world_at(self, env_index: int = None): + + ScenarioUtils.spawn_entities_randomly( + self.world.agents, + self.world, + env_index, + min_dist_between_entities=0, + x_bounds=( + 0, + 0, + ), + y_bounds=( + 0, + 0, + ), + ) + + def reward(self, agent: Agent): + return torch.zeros( + self.world.batch_dim, device=self.world.device, dtype=torch.float32 + ) + + def observation(self, agent: Agent): + return torch.zeros( + self.world.batch_dim, 1, device=self.world.device, dtype=torch.float32 + ) + + +if __name__ == "__main__": + render_interactively(__file__, control_two_agents=True)