Upload 18 files

6953bd1 verified 7 days ago

20.9 kB

"""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])

	"""Deterministic 2D multi-agent navigation simulator.

	Arena 12x12 with walls. Agents are double-integrators:
	v += adt, \|v\| <= VMAX, pos += vdt, \|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])