aparekh02/overflow-openenv
0
1"""2OpenENV RL Demo — Gradio UI entrypoint for HuggingFace Spaces.3 4Runs inside the overflow_env package root. All imports use absolute paths5so they work both as a package (installed) and as a Space (flat root).6"""7 8import sys, os9# When running as HF Space, make server/ importable with absolute paths10sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))11 12import math, time, threading13import numpy as np14import torch15import torch.optim as optim16import matplotlib17matplotlib.use("Agg")18import matplotlib.pyplot as plt19import matplotlib.patches as patches20import gradio as gr21 22from server.overflow_environment import OverflowEnvironment23from models import OverflowAction24from policies.flat_mlp_policy import FlatMLPPolicy25from policies.policy_spec import build_obs, build_ticket_vector, OBS_DIM26 27 28STEPS_PER_EPISODE = 2029NUM_LANES = 330ROAD_LENGTH = 20031 32 33# ── Observation adapter ───────────────────────────────────────────────────────34 35def obs_to_vec(overflow_obs) -> np.ndarray:36 cars = overflow_obs.cars37 if not cars:38 return np.zeros(OBS_DIM, dtype=np.float32)39 ego = next((c for c in cars if c.carId == 0), cars[0])40 ego_spd = ego.speed / 4.541 ego_x = ego.position.x42 ego_y = (ego.lane - 2) * 3.743 tickets = []44 for car in cars:45 if car.carId == 0:46 continue47 rx = car.position.x - ego.position.x48 ry = (car.lane - ego.lane) * 3.749 cs = car.speed / 4.550 d = math.sqrt(rx**2 + ry**2)51 if d > 80:52 continue53 cl = max(ego_spd - cs * math.copysign(1, max(rx, 0.01)), 0.1)54 tickets.append(build_ticket_vector(55 severity_weight=1.0 if d < 8 else 0.75 if d < 15 else 0.5,56 ttl=5.0, pos_x=rx, pos_y=ry, pos_z=0.0,57 vel_x=cs, vel_y=0.0, vel_z=0.0, heading=0.0,58 size_length=4.0, size_width=2.0, size_height=1.5,59 distance=d, time_to_collision=min(d / cl, 30.0),60 bearing=math.atan2(ry, max(rx, 0.01)),61 ticket_type="collision_risk", entity_type="vehicle", confidence=1.0,62 ))63 tv = np.array(tickets, dtype=np.float32) if tickets else None64 return build_obs(ego_x=ego_x, ego_y=ego_y, ego_z=0.0,65 ego_vx=ego_spd, ego_vy=0.0,66 heading=0.0, speed=ego_spd,67 steer=0.0, throttle=0.5, brake=0.0,68 ticket_vectors=tv)69 70 71def action_to_decision(a: np.ndarray) -> str:72 s, t, b = float(a[0]), float(a[1]), float(a[2])73 if abs(s) > 0.35: return "lane_change_left" if s < 0 else "lane_change_right"74 if b > 0.25: return "brake"75 if t > 0.20: return "accelerate"76 return "maintain"77 78 79# ── Global training state ─────────────────────────────────────────────────────80 81policy = FlatMLPPolicy(obs_dim=OBS_DIM)82optimizer = optim.Adam(policy.parameters(), lr=3e-4, eps=1e-5)83 84_buf_obs = []85_buf_acts = []86_buf_rews = []87_buf_logps = []88_buf_vals = []89_buf_dones = []90 91episode_history = []92step_log = []93_running = False94_lock = threading.Lock()95 96 97def _ppo_mini_update():98 if len(_buf_obs) < 2:99 return100 obs_t = torch.tensor(np.array(_buf_obs), dtype=torch.float32)101 acts_t = torch.tensor(np.array(_buf_acts), dtype=torch.float32)102 rews_t = torch.tensor(_buf_rews, dtype=torch.float32)103 logp_t = torch.tensor(_buf_logps, dtype=torch.float32)104 vals_t = torch.tensor(_buf_vals, dtype=torch.float32)105 done_t = torch.tensor(_buf_dones, dtype=torch.float32)106 107 gamma, lam = 0.99, 0.95108 adv = torch.zeros_like(rews_t)109 gae = 0.0110 for t in reversed(range(len(rews_t))):111 nv = 0.0 if t == len(rews_t) - 1 else float(vals_t[t + 1])112 d = rews_t[t] + gamma * nv * (1 - done_t[t]) - vals_t[t]113 gae = d + gamma * lam * (1 - done_t[t]) * gae114 adv[t] = gae115 ret = adv + vals_t116 adv = (adv - adv.mean()) / (adv.std() + 1e-8)117 118 policy.train()119 act_mean, val = policy(obs_t)120 val = val.squeeze(-1)121 dist = torch.distributions.Normal(act_mean, torch.ones_like(act_mean) * 0.3)122 logp = dist.log_prob(acts_t).sum(dim=-1)123 entropy = dist.entropy().sum(dim=-1).mean()124 ratio = torch.exp(logp - logp_t)125 pg = torch.max(-adv * ratio, -adv * ratio.clamp(0.8, 1.2)).mean()126 vf = 0.5 * ((val - ret) ** 2).mean()127 loss = pg + 0.5 * vf - 0.02 * entropy128 optimizer.zero_grad()129 loss.backward()130 torch.nn.utils.clip_grad_norm_(policy.parameters(), 0.5)131 optimizer.step()132 133 134def run_episodes_loop():135 global _running136 ep_num = 0137 env = OverflowEnvironment()138 139 while _running:140 ep_num += 1141 obs = env.reset()142 ep_rew = 0.0143 outcome = "timeout"144 145 _buf_obs.clear(); _buf_acts.clear(); _buf_rews.clear()146 _buf_logps.clear(); _buf_vals.clear(); _buf_dones.clear()147 148 for step in range(1, STEPS_PER_EPISODE + 1):149 if not _running:150 break151 152 obs_vec = obs_to_vec(obs)153 policy.eval()154 with torch.no_grad():155 obs_t = torch.tensor(obs_vec, dtype=torch.float32).unsqueeze(0)156 act_mean, val = policy(obs_t)157 dist = torch.distributions.Normal(act_mean.squeeze(0),158 torch.ones(3) * 0.3)159 action = dist.sample().clamp(-1, 1)160 logp = dist.log_prob(action).sum()161 162 decision = action_to_decision(action.numpy())163 obs = env.step(OverflowAction(decision=decision, reasoning=""))164 reward = float(obs.reward or 0.0)165 done = obs.done166 ep_rew += reward167 168 _buf_obs.append(obs_vec)169 _buf_acts.append(action.numpy())170 _buf_rews.append(reward)171 _buf_logps.append(float(logp))172 _buf_vals.append(float(val.squeeze()))173 _buf_dones.append(float(done))174 175 with _lock:176 step_log.append({177 "ep": ep_num,178 "step": step,179 "decision": decision,180 "reward": round(reward, 2),181 "ep_reward": round(ep_rew, 2),182 "incident": obs.incident_report or "",183 "cars": [(c.carId, c.lane, c.position.x, c.speed)184 for c in obs.cars],185 })186 187 if done:188 outcome = "CRASH" if "CRASH" in (obs.incident_report or "") else "GOAL"189 break190 191 time.sleep(0.6)192 193 _ppo_mini_update()194 195 with _lock:196 episode_history.append({197 "ep": ep_num,198 "steps": step,199 "reward": round(ep_rew, 2),200 "outcome": outcome,201 })202 203 204# ── Plot helpers ──────────────────────────────────────────────────────────────205 206DECISION_COLORS = {207 "accelerate": "#22c55e",208 "brake": "#ef4444",209 "lane_change_left": "#f59e0b",210 "lane_change_right": "#f59e0b",211 "maintain": "#60a5fa",212}213 214 215def render_road(cars_snapshot, last_decision, last_incident):216 fig, ax = plt.subplots(figsize=(10, 2.8))217 fig.patch.set_facecolor("#0f172a")218 ax.set_facecolor("#1e293b")219 220 ax.set_xlim(0, ROAD_LENGTH)221 ax.set_ylim(0, NUM_LANES + 1)222 ax.set_yticks([])223 ax.set_xlabel("Position", color="#94a3b8", fontsize=9)224 ax.tick_params(colors="#94a3b8")225 for spine in ax.spines.values():226 spine.set_edgecolor("#334155")227 228 for lane in range(1, NUM_LANES):229 ax.axhline(y=lane + 0.5, color="#334155", linewidth=1, linestyle="--", alpha=0.6)230 231 for lane in range(1, NUM_LANES + 1):232 ax.text(2, lane, f"L{lane}", color="#475569", fontsize=8, va="center")233 234 ax.axvspan(160, ROAD_LENGTH, alpha=0.12, color="#22c55e")235 ax.text(162, NUM_LANES + 0.6, "GOAL ZONE", color="#22c55e", fontsize=7, alpha=0.8)236 237 car_w, car_h = 8, 0.55238 for car_id, lane, pos_x, speed in cars_snapshot:239 is_ego = car_id == 0240 color = "#3b82f6" if is_ego else "#94a3b8"241 outline = "#60a5fa" if is_ego else "#475569"242 lw = 2.0 if is_ego else 1.0243 rect = patches.FancyBboxPatch(244 (pos_x - car_w / 2, lane - car_h / 2),245 car_w, car_h,246 boxstyle="round,pad=0.05",247 facecolor=color, edgecolor=outline, linewidth=lw, alpha=0.92,248 )249 ax.add_patch(rect)250 label = f"{'EGO' if is_ego else f'C{car_id}'}\n{speed:.0f}"251 ax.text(pos_x, lane, label, ha="center", va="center",252 fontsize=6.5, color="white", fontweight="bold" if is_ego else "normal")253 254 dec_color = DECISION_COLORS.get(last_decision, "#60a5fa")255 ax.text(ROAD_LENGTH - 2, NUM_LANES + 0.65,256 f"Action: {last_decision.replace('_', ' ').upper()}",257 color=dec_color, fontsize=8, fontweight="bold", ha="right")258 259 if "CRASH" in last_incident:260 ax.text(ROAD_LENGTH / 2, NUM_LANES + 0.65, "CRASH",261 color="#ef4444", fontsize=10, fontweight="bold", ha="center")262 elif "NEAR MISS" in last_incident:263 ax.text(ROAD_LENGTH / 2, NUM_LANES + 0.65, "NEAR MISS",264 color="#f59e0b", fontsize=9, fontweight="bold", ha="center")265 elif "GOAL" in last_incident:266 ax.text(ROAD_LENGTH / 2, NUM_LANES + 0.65, "GOAL REACHED",267 color="#22c55e", fontsize=10, fontweight="bold", ha="center")268 269 plt.tight_layout(pad=0.3)270 return fig271 272 273def render_reward_curve(eps):274 fig, ax = plt.subplots(figsize=(10, 2.8))275 fig.patch.set_facecolor("#0f172a")276 ax.set_facecolor("#1e293b")277 for spine in ax.spines.values():278 spine.set_edgecolor("#334155")279 ax.tick_params(colors="#94a3b8")280 ax.set_xlabel("Episode", color="#94a3b8", fontsize=9)281 ax.set_ylabel("Total Reward", color="#94a3b8", fontsize=9)282 283 if not eps:284 ax.text(0.5, 0.5, "Waiting for episodes...", transform=ax.transAxes,285 ha="center", va="center", color="#475569", fontsize=11)286 plt.tight_layout(pad=0.3)287 return fig288 289 xs = [e["ep"] for e in eps]290 ys = [e["reward"] for e in eps]291 outcome_colors = {"CRASH": "#ef4444", "GOAL": "#22c55e", "timeout": "#60a5fa"}292 for x, y, e in zip(xs, ys, eps):293 ax.bar(x, y, color=outcome_colors.get(e["outcome"], "#60a5fa"), alpha=0.6, width=0.7)294 295 if len(ys) >= 3:296 w = min(5, len(ys))297 smoothed = np.convolve(ys, np.ones(w) / w, mode="valid")298 ax.plot(xs[w - 1:], smoothed, color="#f8fafc", linewidth=2)299 300 ax.axhline(0, color="#334155", linewidth=0.8)301 302 from matplotlib.patches import Patch303 legend_els = [Patch(facecolor="#ef4444", label="crash"),304 Patch(facecolor="#22c55e", label="goal"),305 Patch(facecolor="#60a5fa", label="timeout")]306 ax.legend(handles=legend_els, facecolor="#1e293b", labelcolor="#94a3b8",307 fontsize=8, framealpha=0.6, edgecolor="#334155", loc="upper left")308 309 plt.tight_layout(pad=0.3)310 return fig311 312 313# ── Gradio UI ─────────────────────────────────────────────────────────────────314 315def start_training():316 global _running317 if not _running:318 _running = True319 step_log.clear()320 episode_history.clear()321 threading.Thread(target=run_episodes_loop, daemon=True).start()322 return gr.update(value="Running...", interactive=False), gr.update(interactive=True)323 324 325def stop_training():326 global _running327 _running = False328 return gr.update(value="Start", interactive=True), gr.update(interactive=False)329 330 331def get_updates():332 with _lock:333 logs = list(step_log[-20:])334 eps = list(episode_history[-50:])335 last = step_log[-1] if step_log else None336 337 road_fig = render_road(last["cars"], last["decision"], last["incident"]) if last \338 else render_road([], "maintain", "")339 reward_fig = render_reward_curve(eps)340 341 lines = []342 for e in reversed(logs):343 flag = ""344 if "CRASH" in e["incident"]: flag = " 💥"345 elif "GOAL" in e["incident"]: flag = " ✓"346 elif "NEAR MISS" in e["incident"]: flag = " ⚠"347 lines.append(348 f"ep {e['ep']:>3d} | step {e['step']:>2d} | "349 f"{e['decision']:<20} | r={e['reward']:>+6.2f} | "350 f"ep_total={e['ep_reward']:>7.2f}{flag}"351 )352 step_text = "\n".join(lines) if lines else "Waiting for first episode..."353 354 ep_lines = ["Episode | Steps | Total Reward | Outcome", "-" * 44]355 for e in reversed(eps[-15:]):356 ep_lines.append(357 f" {e['ep']:>4d} | {e['steps']:>3d} | "358 f" {e['reward']:>+8.2f} | {e['outcome']}"359 )360 ep_text = "\n".join(ep_lines) if eps else "No episodes completed yet."361 362 if len(eps) >= 2:363 rewards = [e["reward"] for e in eps]364 n = len(rewards)365 half = max(n // 2, 1)366 early = sum(rewards[:half]) / half367 late = sum(rewards[half:]) / max(n - half, 1)368 arrow = "↑ improving" if late > early else "↓ declining"369 trend_text = f"Early {half} eps: {early:+.2f} → Last {n-half} eps: {late:+.2f} {arrow}"370 else:371 trend_text = "Collecting data..."372 373 status = "● RUNNING" if _running else "■ STOPPED"374 return road_fig, reward_fig, step_text, ep_text, trend_text, status375 376 377_EMPTY_ROAD = render_road([], "maintain", "")378_EMPTY_REWARD = render_reward_curve([])379 380with gr.Blocks(title="OpenENV RL Demo", theme=gr.themes.Base()) as demo:381 gr.Markdown(382 "# OpenENV RL — Live Policy Training\n"383 "**FlatMLPPolicy** drives Car 0 on a 3-lane road for 20 steps per episode. "384 "PPO mini-update after each episode — watch rewards trend upward over time."385 )386 387 with gr.Row():388 start_btn = gr.Button("Start", variant="primary", scale=1)389 stop_btn = gr.Button("Stop", variant="stop", interactive=False, scale=1)390 status_box = gr.Textbox(value="■ STOPPED", label="Status",391 interactive=False, scale=0, min_width=130)392 393 gr.Markdown("### Road View")394 road_plot = gr.Plot(value=_EMPTY_ROAD, show_label=False)395 396 gr.Markdown("### Episode Reward Curve")397 reward_plot = gr.Plot(value=_EMPTY_REWARD, show_label=False)398 399 gr.Markdown("### Live Step Feed (last 20 steps)")400 step_display = gr.Textbox(401 value="Press Start to begin...",402 lines=14, max_lines=14, interactive=False,403 )404 405 with gr.Row():406 with gr.Column():407 gr.Markdown("### Episode History")408 ep_display = gr.Textbox(lines=10, interactive=False)409 with gr.Column():410 gr.Markdown("### Reward Trend")411 trend_display = gr.Textbox(lines=3, interactive=False)412 413 timer = gr.Timer(value=1.0)414 timer.tick(415 fn=get_updates,416 outputs=[road_plot, reward_plot, step_display, ep_display, trend_display, status_box],417 )418 419 start_btn.click(fn=start_training, outputs=[start_btn, stop_btn])420 stop_btn.click(fn=stop_training, outputs=[start_btn, stop_btn])421 422 423if __name__ == "__main__":424 demo.launch()425