"""Deterministic 2D multi-agent navigation simulator. Arena 12x12 with walls. Agents are double-integrators: v += a*dt, |v| <= VMAX, pos += v*dt, |a| <= AMAX, dt = 0.1. Agent radius 0.25. Collision with a circle obstacle, a wall segment, an arena wall, or another agent ends that agent's episode as a failure. Success = goal reached (dist < 0.4) with no collision in the episode. Observations are decentralized and shared-policy: each agent sees only its own velocity, relative goal, nearest circle obstacle, nearest wall segment, nearest other agent, and nearest arena-wall distance. OBS layout (15 dims, see OBS_NAMES): 0,1 own velocity / VMAX 2,3 relative goal vector / 12 4 goal distance / 12 5,6 unit vector away from nearest circle obstacle (surface point) 7 surface distance to nearest circle obstacle / 12 8,9 unit vector away from nearest wall-segment point 10 distance to nearest wall segment / 12 11,12 unit vector away from nearest other agent 13 centre distance to nearest other agent / 12 (1.0 if no other agent) 14 distance to nearest arena wall / 12 15,16 relative vector to wall-segment endpoint 1 / 12 17,18 relative vector to wall-segment endpoint 2 / 12 19 radius of nearest circle obstacle (0.5-0.8) Dims 15-19 were added after a failed training run: with only the nearest segment point, the policy could not observe where the wall ends, so choosing a gap was unobservable and the wall-trap detour was unlearnable. Endpoint vectors are standard line-feature sensing (e.g. a lidar return's endpoints). They are appended at the end so the classical controller's indices 0-14 are unchanged. The batched core (BatchedSim) is the single source of physics; NavSim wraps it with n_envs=1 and records trajectories for plotting/eval. """ import numpy as np ARENA = 12.0 DT = 0.1 # Limits chosen so the control problem is learnable: at VMAX 2.0 the # stopping distance (v^2/2a = 0.5 m) barely fit inside the 0.8 m # safety zone, so exploration noise alone collided near hazards. # At 1.2/3.0 the stopping distance is 0.24 m. Changed before the # successful training run; classical gains were re-checked after. VMAX = 1.2 AMAX = 3.0 AGENT_R = 0.25 GOAL_TOL = 0.4 MAX_STEPS = 200 OBS_DIM = 20 OBS_NAMES = [ "vel_x", "vel_y", "goal_dx", "goal_dy", "goal_dist", "obs_dir_x", "obs_dir_y", "obs_dist", "seg_dir_x", "seg_dir_y", "seg_dist", "neigh_dir_x", "neigh_dir_y", "neigh_dist", "wall_dist", "seg_end1_dx", "seg_end1_dy", "seg_end2_dx", "seg_end2_dy", "obs_radius", ] SAFETY_MARGIN = 0.15 # clearance below this counts as a safety violation def _seg_closest(px, py, x1, y1, x2, y2): """Closest point on segment to (px,py); vectorised over leading axes.""" dx, dy = x2 - x1, y2 - y1 L2 = dx * dx + dy * dy t = np.where(L2 > 0, ((px - x1) * dx + (py - y1) * dy) / np.maximum(L2, 1e-9), 0.0) t = np.clip(t, 0.0, 1.0) cx, cy = x1 + t * dx, y1 + t * dy return cx, cy def generate_layout(rng, n_agents, trap=None): """Seeded layout. trap=True forces the vertical wall-trap layout. Trap: vertical segment x=6, y in [2,10]; gaps only at y<2 and y>10. Starts on the left, goals on the right, so a direct radial potential field is attracted right and repelled left by the wall in the middle. """ if trap is None: trap = bool(rng.random() < 0.5) attempts = 0 while True: attempts += 1 if attempts > 500: raise RuntimeError("layout generation failed") if trap: segments = np.array([[6.0, 2.0, 6.0, 10.0]], dtype=np.float64) # Goals near-directly opposite the starts, mid-wall, so the # direct route meets the wall head-on: the textbook local # minimum. (A large y-offset would let the field slide # along the wall to a gap and would not be a trap at all.) ys = np.linspace(4.5, 7.5, n_agents) + rng.uniform(-0.3, 0.3, size=n_agents) ys = np.clip(ys, 4.0, 8.0) starts = np.stack([rng.uniform(1.0, 3.0, size=n_agents), ys], axis=1) perm = rng.permutation(n_agents) goal_ys = ys[perm] + rng.uniform(-0.4, 0.4, size=n_agents) goals = np.stack([rng.uniform(9.0, 11.0, size=n_agents), np.clip(goal_ys, 1.2, 10.8)], axis=1) circles = [] while len(circles) < 3: cx = rng.uniform(0.8, 11.2) # keep circles off the wall-trap corridor so the wall # itself, not a circle, is what blocks the direct route if 4.8 < cx < 7.2: continue cy = rng.uniform(0.8, 11.2) r = rng.uniform(0.5, 0.8) cand = np.array([cx, cy, r]) if _layout_ok(starts, goals, circles, cand): circles.append(cand) circles = np.array(circles, dtype=np.float64) else: # one shorter barrier segment, placed mid-arena at a random # angle, plus 3 circles anywhere (with separation checks) ang = rng.uniform(0, np.pi) cx0, cy0 = rng.uniform(4.0, 8.0), rng.uniform(4.0, 8.0) half = 1.5 segments = np.array([[ cx0 - half * np.cos(ang), cy0 - half * np.sin(ang), cx0 + half * np.cos(ang), cy0 + half * np.sin(ang), ]], dtype=np.float64) starts = _spread_points(rng, n_agents, lo=0.8, hi=11.2) goals = _spread_points(rng, n_agents, lo=0.8, hi=11.2) # matched start->goal distances must be meaningful trips if np.min(np.linalg.norm(starts - goals, axis=1)) < 4.0: continue # starts/goals must not spawn on/inside the barrier segment pts = np.vstack([starts, goals]) cx, cy = _seg_closest(pts[:, 0], pts[:, 1], segments[0, 0], segments[0, 1], segments[0, 2], segments[0, 3]) if np.min(np.hypot(pts[:, 0] - cx, pts[:, 1] - cy)) < 1.0: continue circles = [] while len(circles) < 3: cand = np.array([rng.uniform(0.8, 11.2), rng.uniform(0.8, 11.2), rng.uniform(0.5, 0.8)]) if _layout_ok(starts, goals, circles, cand, seg=segments[0]): circles.append(cand) circles = np.array(circles, dtype=np.float64) return {"starts": starts, "goals": goals, "circles": circles, "segments": segments, "trap": bool(trap)} def _spread_points(rng, n, lo, hi): pts = [] while len(pts) < n: cand = rng.uniform(lo, hi, size=2) if all(np.linalg.norm(cand - p) > 1.2 for p in pts): pts.append(cand) return np.array(pts, dtype=np.float64) def _layout_ok(starts, goals, circles, cand, seg=None): for p in np.vstack([starts, goals]): if np.linalg.norm(cand[:2] - p) < cand[2] + 0.9: return False for c in circles: if np.linalg.norm(cand[:2] - c[:2]) < cand[2] + c[2] + 0.6: return False if seg is not None: cx, cy = _seg_closest(cand[0], cand[1], seg[0], seg[1], seg[2], seg[3]) if np.hypot(cand[0] - cx, cand[1] - cy) < cand[2] + 0.5: return False return True class BatchedSim: """Vectorised (n_envs, n_agents) simulator. All physics lives here.""" def __init__(self, n_envs, n_agents, seed=0): self.E, self.A = n_envs, n_agents self.rng = np.random.default_rng(seed) self.pos = np.zeros((n_envs, n_agents, 2)) self.vel = np.zeros((n_envs, n_agents, 2)) self.goals = np.zeros((n_envs, n_agents, 2)) self.circles = np.zeros((n_envs, 3, 3)) self.segments = np.zeros((n_envs, 1, 4)) self.active = np.ones((n_envs, n_agents), dtype=bool) self.done_agent = np.zeros((n_envs, n_agents), dtype=bool) self.success = np.zeros((n_envs, n_agents), dtype=bool) self.collision = np.zeros((n_envs, n_agents), dtype=bool) self.collision_kind = [["" for _ in range(n_agents)] for _ in range(n_envs)] self.steps = np.zeros(n_envs, dtype=int) self.prev_goal_dist = np.zeros((n_envs, n_agents)) self.path_len = np.zeros((n_envs, n_agents)) self.min_clearance = np.full((n_envs, n_agents), np.inf) self.viol_steps = np.zeros((n_envs, n_agents), dtype=int) self.active_steps = np.zeros((n_envs, n_agents), dtype=int) # ---- layout / reset ------------------------------------------------- def reset(self, layouts, obs_noise=0.0, disturbance=0.0): self.obs_noise = float(obs_noise) self.disturbance = float(disturbance) for e, lay in enumerate(layouts): self.pos[e] = lay["starts"] self.goals[e] = lay["goals"] self.circles[e] = lay["circles"] self.segments[e] = lay["segments"] self.vel[:] = 0.0 self.active[:] = True self.done_agent[:] = False self.success[:] = False self.collision[:] = False self.collision_kind = [["" for _ in range(self.A)] for _ in range(self.E)] self.steps[:] = 0 self.path_len[:] = 0.0 self.min_clearance[:] = np.inf self.viol_steps[:] = 0 self.active_steps[:] = 0 self.prev_goal_dist = np.linalg.norm(self.goals - self.pos, axis=2) self._update_clearance() return self._obs() def reset_envs(self, idxs, layouts): """Reset only the listed envs (used by training partial resets).""" for e, lay in zip(idxs, layouts): self.pos[e] = lay["starts"] self.goals[e] = lay["goals"] self.circles[e] = lay["circles"] self.segments[e] = lay["segments"] self.vel[e] = 0.0 self.active[e] = True self.done_agent[e] = False self.success[e] = False self.collision[e] = False self.collision_kind[e] = ["" for _ in range(self.A)] self.steps[e] = 0 self.path_len[e] = 0.0 self.min_clearance[e] = np.inf self.viol_steps[e] = 0 self.active_steps[e] = 0 self.prev_goal_dist[e] = np.linalg.norm(self.goals[e] - self.pos[e], axis=1) cl = self._clearance() self.min_clearance[idxs] = cl[idxs] self.viol_steps[idxs] += (cl[idxs] < SAFETY_MARGIN).astype(int) return self._obs() # ---- geometry helpers ---------------------------------------------- def _clearance(self): """Per-agent clearance: min(surface distance) to any hazard.""" E, A = self.E, self.A # circles: centre dist - r_circle - AGENT_R d_c = np.linalg.norm(self.pos[:, :, None, :] - self.circles[:, None, :, :2], axis=3) \ - self.circles[:, None, :, 2] - AGENT_R # (E,A,C) # segments: point-segment dist - AGENT_R seg = self.segments # (E,1,4) px = self.pos[:, :, None, 0] py = self.pos[:, :, None, 1] cx, cy = _seg_closest(px, py, seg[:, :, 0][:, :, None], seg[:, :, 1][:, :, None], seg[:, :, 2][:, :, None], seg[:, :, 3][:, :, None]) d_s = np.hypot(px - cx, py - cy) - AGENT_R # (E,A,1) # arena walls d_w = np.minimum(np.minimum(self.pos[..., 0], ARENA - self.pos[..., 0]), np.minimum(self.pos[..., 1], ARENA - self.pos[..., 1])) - AGENT_R # other agents: centre dist - 2*AGENT_R (ignore done agents) d_pair = np.linalg.norm(self.pos[:, :, None, :] - self.pos[:, None, :, :], axis=3) - 2 * AGENT_R eye = np.eye(A, dtype=bool) d_pair[:, eye] = np.inf d_pair = np.where(self.active[:, None, :] | self.active[:, :, None], d_pair, np.inf) # done agents are not hazards for the remaining agents d_pair = np.where(self.done_agent[:, None, :], np.inf, d_pair) d_pair[:, eye] = np.inf return np.minimum(np.minimum(d_c.min(axis=2), d_s[:, :, 0]), np.minimum(d_w, d_pair.min(axis=2))) def _update_clearance(self): cl = self._clearance() act = self.active self.min_clearance = np.where(act, np.minimum(self.min_clearance, cl), self.min_clearance) self.viol_steps += np.where(act & (cl < SAFETY_MARGIN), 1, 0) def _obs(self): E, A = self.E, self.A obs = np.zeros((E, A, OBS_DIM), dtype=np.float64) obs[:, :, 0:2] = self.vel / VMAX rel = self.goals - self.pos obs[:, :, 2:4] = rel / ARENA obs[:, :, 4] = np.linalg.norm(rel, axis=2) / ARENA # nearest circle (surface point) d_c = np.linalg.norm(self.pos[:, :, None, :] - self.circles[:, None, :, :2], axis=3) # centre dist idx = np.argmin(d_c - self.circles[:, None, :, 2], axis=2) # nearest by surface dist surf = np.take_along_axis(d_c, idx[:, :, None], axis=2)[:, :, 0] \ - self.circles[np.arange(E)[:, None], idx, 2] centres = self.circles[np.arange(E)[:, None], idx][:, :, :2] # (E,A,2) away = self.pos - centres nrm = np.maximum(np.linalg.norm(away, axis=2, keepdims=True), 1e-6) obs[:, :, 5:7] = away / nrm obs[:, :, 7] = np.clip(surf, 0.0, ARENA) / ARENA # nearest segment point seg = self.segments px = self.pos[:, :, None, 0] py = self.pos[:, :, None, 1] cx, cy = _seg_closest(px, py, seg[:, :, 0][:, :, None], seg[:, :, 1][:, :, None], seg[:, :, 2][:, :, None], seg[:, :, 3][:, :, None]) d_seg = np.hypot(px - cx, py - cy)[:, :, 0] away_s = np.stack([self.pos[:, :, 0] - cx[:, :, 0], self.pos[:, :, 1] - cy[:, :, 0]], axis=2) nrm_s = np.maximum(np.linalg.norm(away_s, axis=2, keepdims=True), 1e-6) obs[:, :, 8:10] = away_s / nrm_s obs[:, :, 10] = d_seg / ARENA # nearest other agent d_pair = np.linalg.norm(self.pos[:, :, None, :] - self.pos[:, None, :, :], axis=3) d_pair[:, np.eye(A, dtype=bool)] = np.inf d_pair = np.where(self.done_agent[:, None, :], np.inf, d_pair) d_pair[:, np.eye(A, dtype=bool)] = np.inf j = np.argmin(d_pair, axis=2) neigh = self.pos[np.arange(E)[:, None], j] # (E,A,2) away_n = self.pos - neigh nrm_n = np.maximum(np.linalg.norm(away_n, axis=2, keepdims=True), 1e-6) has_neigh = np.isfinite(d_pair.min(axis=2)) obs[:, :, 11:13] = np.where(has_neigh[:, :, None], away_n / nrm_n, 0.0) obs[:, :, 13] = np.where(has_neigh, np.minimum(d_pair.min(axis=2), ARENA), ARENA) / ARENA # arena wall distance obs[:, :, 14] = np.minimum(np.minimum(self.pos[..., 0], ARENA - self.pos[..., 0]), np.minimum(self.pos[..., 1], ARENA - self.pos[..., 1])) / ARENA # segment endpoints (relative vectors) + nearest circle radius obs[:, :, 15:17] = (seg[:, :, :2] - self.pos) / ARENA obs[:, :, 17:19] = (seg[:, :, 2:4] - self.pos) / ARENA obs[:, :, 19] = self.circles[np.arange(E)[:, None], idx, 2] if self.obs_noise > 0: obs = obs + self.rng.normal(0.0, self.obs_noise, size=obs.shape) obs[:, :, 0:2] = np.clip(obs[:, :, 0:2], -1.5, 1.5) return obs.astype(np.float32) # ---- step ------------------------------------------------------------ def step(self, actions): """actions: (E,A,2) accelerations, clipped to AMAX. Returns obs, reward, episode_done (E,), info dict. Per-agent done flags are in self.done_agent; reward is the training shaping reward.""" actions = np.clip(np.asarray(actions, dtype=np.float64), -AMAX, AMAX) if self.disturbance > 0: actions = actions + self.rng.normal(0.0, self.disturbance, size=actions.shape) actions = np.clip(actions, -AMAX, AMAX) act = self.active old_pos = self.pos.copy() self.vel = np.where(act[:, :, None], self.vel + actions * DT, self.vel) spd = np.linalg.norm(self.vel, axis=2, keepdims=True) self.vel = np.where(spd > VMAX, self.vel * (VMAX / np.maximum(spd, 1e-9)), self.vel) self.pos = np.where(act[:, :, None], self.pos + self.vel * DT, self.pos) self.path_len += np.where(act, np.linalg.norm(self.pos - old_pos, axis=2), 0.0) self.active_steps += act.astype(int) self.steps += 1 cl = self._clearance() self.min_clearance = np.where(act, np.minimum(self.min_clearance, cl), self.min_clearance) self.viol_steps += np.where(act & (cl < SAFETY_MARGIN), 1, 0) goal_dist = np.linalg.norm(self.goals - self.pos, axis=2) progress = self.prev_goal_dist - goal_dist # Proximity shaping. A stronger 2.0/0.8 version was tried and # made learning worse (returns ~-12, success ~10%): agents were # punished heavily just for starting near hazards. Reverted to # the original milder term; the curriculum (trap ramp in # train.py) is what teaches avoidance instead. reward = 2.0 * progress - 0.02 - 0.5 * np.maximum(0.0, 0.4 - np.where(act, cl, 0.4)) reward = np.where(act, reward, 0.0) # collisions (checked per kind for reporting) d_c = np.linalg.norm(self.pos[:, :, None, :] - self.circles[:, None, :, :2], axis=3) \ - self.circles[:, None, :, 2] hit_circle = (d_c < AGENT_R).any(axis=2) seg = self.segments px = self.pos[:, :, None, 0] py = self.pos[:, :, None, 1] cx, cy = _seg_closest(px, py, seg[:, :, 0][:, :, None], seg[:, :, 1][:, :, None], seg[:, :, 2][:, :, None], seg[:, :, 3][:, :, None]) hit_seg = (np.hypot(px - cx, py - cy)[:, :, 0] < AGENT_R) hit_wall = (self.pos[..., 0] < AGENT_R) | (self.pos[..., 0] > ARENA - AGENT_R) | \ (self.pos[..., 1] < AGENT_R) | (self.pos[..., 1] > ARENA - AGENT_R) d_pair = np.linalg.norm(self.pos[:, :, None, :] - self.pos[:, None, :, :], axis=3) d_pair[:, np.eye(self.A, dtype=bool)] = np.inf d_pair = np.where(self.done_agent[:, None, :], np.inf, d_pair) d_pair[:, np.eye(self.A, dtype=bool)] = np.inf hit_agent = (d_pair < 2 * AGENT_R).any(axis=2) hit_obstacle = act & (hit_circle | hit_seg | hit_wall) hit_other = act & hit_agent reached = act & (goal_dist < GOAL_TOL) # collision takes precedence over reaching on the same step newly_coll = act & (hit_obstacle | hit_other) newly_ok = reached & ~newly_coll for e in range(self.E): for a in range(self.A): if not act[e, a]: continue if hit_obstacle[e, a] and hit_other[e, a]: self.collision_kind[e][a] = "obstacle+agent" elif hit_obstacle[e, a]: self.collision_kind[e][a] = "obstacle" elif hit_other[e, a]: self.collision_kind[e][a] = "agent" self.collision |= newly_coll self.success |= newly_ok reward = np.where(newly_coll, reward - 10.0, reward) reward = np.where(newly_ok, reward + 10.0, reward) self.done_agent |= newly_coll | newly_ok self.active = ~self.done_agent self.prev_goal_dist = goal_dist timed_out = (self.steps >= MAX_STEPS) ep_done = np.all(self.done_agent, axis=1) | timed_out # agents still active at timeout are failures (timeout, not collision) if timed_out.any(): self.active[timed_out] = False return self._obs(), reward.astype(np.float32), ep_done, {} class NavSim: """Single-episode wrapper recording trajectories (for eval/plots).""" def __init__(self, n_agents, seed=0): self.core = BatchedSim(1, n_agents, seed=seed) self.n_agents = n_agents self.traj = None self.layout = None def reset(self, layout, obs_noise=0.0, disturbance=0.0): self.layout = layout self.traj = [ [layout["starts"][i].copy()] for i in range(self.n_agents) ] return self.core.reset([layout], obs_noise=obs_noise, disturbance=disturbance)[0] def step(self, action): obs, rew, done, _ = self.core.step(action[None, :, :]) for i in range(self.n_agents): self.traj[i].append(self.core.pos[0, i].copy()) return obs[0], rew[0], bool(done[0]), {} @property def pos(self): return self.core.pos[0] @property def success(self): return self.core.success[0] @property def collision(self): return self.core.collision[0] @property def collision_kind(self): return self.core.collision_kind[0] @property def steps(self): return int(self.core.steps[0])