🎮 Gymnasium RL Environment

Train reinforcement learning policies on city-scale traffic using GPU simulation as the environment

Architecture

🤖
RL Agent
PPO / SAC / MoE
🔄
Gym Interface
obs → action → reward
GPU Simulator
223K nodes, 540K edges
Agent ←→ Environment ←→ CUDA Kernels

Live Demo Output (from actual GPU run)

Python 3.10Gymnasium 1.3NumPy
>>> from lpsim_env import LPSimEnv

>>> env = LPSimEnv( ... network_path="data/networks/sf_bay_area", ... reward_type="travel_time", ... num_trips=1000 ... ) Action space: Box(low=0, high=1, shape=(100,)) Observation space: Box(shape=(547697, 3)) — [speed, density, flow] per edge Network: 547,697 edges, 100 controllable

>>> obs, info = env.reset(seed=42) obs.shape = (547697, 3) obs.dtype = float32

>>> action = env.action_space.sample() action.shape = (100,) action[:5] = [0.050, 0.327, 0.650, 0.845, 0.252]

>>> obs, reward, terminated, truncated, info = env.step(action) # GPU simulator executes: 223K nodes, 540K edges, 702 vehicles reward = -26.02 # negative avg travel time (minimize) terminated = True info = { 'avg_travel_time': 26.02, # minutes 'total_co': 470.81, # CO emissions (grams) 'num_completed': 702 # trips finished in window }

>>> # Train with stable-baselines3 >>> from stable_baselines3 import PPO >>> model = PPO('MlpPolicy', env, verbose=1) >>> model.learn(total_timesteps=10000)

Reward Functions

# Available reward types: env = LPSimEnv(reward_type="travel_time") # minimize avg travel time env = LPSimEnv(reward_type="emissions") # minimize CO emissions env = LPSimEnv(reward_type="throughput") # maximize completed trips

# Custom reward: class MyEnv(LPSimEnv): def _compute_reward(self, results): return -results['avg_travel_time'] * 0.7 - results['total_co'] * 0.3

Action Space

# Default: toll vector on 100 highest-traffic edges # action[i] ∈ [0, 1] maps to toll amount $0-$10

# Customize which edges are controllable: env = LPSimEnv( action_edges=[12345, 67890, 11111], # specific edge IDs num_trips=5000 )

# Full network control (540K actions — use with hierarchical RL): env = LPSimEnv(action_edges=list(range(env.num_edges)))