File size: 20,919 Bytes
6953bd1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 | """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])
|