cpsat-scenario-planner

A behavioral planner for highway scenarios using OR-Tools CP-SAT. Intended as a research and education codebase — small enough to read in an afternoon.

---

How It Works

The planner runs a receding-horizon MPC loop: at each simulation step it builds a CP-SAT model, solves for a feasible plan over the next H steps, then executes only the first action and repeats.

for each sim step:
  1. advance oracles          — observed agents move forward (step_fn)
  2. predict oracles          — constant-velocity forecast for H steps (predict_fn)
  3. solve CP model           — find feasible (seg, lane, speed) for planned agents
  4. execute action[0]        — only the first step is committed
  5. check goals              — mark done, stop early if all terminal goals met

Decision variables (integer, per planned agent per horizon step):

Variable	Meaning
seg[t]	longitudinal position [solver segments]
lane[t]	lane index
step[t]	speed [segs/step]
lat[t]	lateral change this step (−1, 0, +1)
is_changing[t]	bool — currently executing a lane change

Constraints encoded in the CP model:

• Kinematic: speed in [min_step, max_step], acceleration ≤ max_step_delta per step
• Safety: minimum longitudinal gap to every other agent in the same lane; larger gap during lane changes (both source and destination lane)
• Lane change duration: once a lane change starts, lat = 0 for the next change_dur − 1 steps
• Inter-agent: same gap rules enforced between pairs of planned agents
• Goals: terminal goals encoded as AddElement constraints (must be satisfied at some t_g ∈ [1, H]); permanent goals enforced as hard constraints at every horizon step

Objective: minimize Σ t_g (reach goals early) + jerk + lane change count. All goal and safety constraints are hard — there is no soft fallback.

---

Limitations

• Discrete state space. Position and speed are integers. The Frenet spline is fitted post-hoc from the waypoints; the solver never sees continuous dynamics or vehicle geometry.
• 1D road model. Lanes are integers, positions are 1D projections. No road curvature, intersections, or topology.
• Prediction mismatch by design. The planner predicts oracles at constant velocity regardless of their actual step_fn. This is intentional (mirrors a real planner's limited knowledge), but it means the planner can be surprised.
• Hard constraints only. If the goal combination is infeasible within the horizon, the solver raises InfeasibleScenarioError instead of finding the best compromise. Tune horizon and max_time_s if you hit this.
• Solve time is not guaranteed. max_time_s is a budget, not a deadline. CP-SAT may return a feasible (non-optimal) solution, or time out entirely.
• Sequential goals. Goal i is only attempted once goal i−1 is done. There is no way to express concurrent or conditional goal structure.
• No explicit cost function tuning. The objective weights (jerk, lane change, goal timing) are fixed. Changing them requires editing mpc_step().

---

Architecture

solver.py        — IR types + CP-SAT MPC engine
pipeline.py      — assembles SimResult from raw solver output
trajectory.py    — Frenet spline from discrete waypoints
visualization.py — bird's-eye animation + metrics plot
scenarios/       — ready-to-run examples

---

Quickstart

make install   # install dependencies

make elk       # ELK overtaking  — 1 planned agent, 1 oracle
make merge     # highway merge   — 2 planned agents, 1 oracle
make overtake  # joint overtake  — 2 planned agents, 3 oracles

make test
make lint

Minimal example:

from cpsat_scenario_planner.solver import (
    CarState, ConstraintSet, Goal, OracleSpec,
    ScenarioDefinition, SolverParams, WorldModel,
    make_constant_oracle, run_mpc_simulation,
)
from cpsat_scenario_planner.pipeline import VehicleDims, build_sim_result
from cpsat_scenario_planner.visualization import animate, plot_metrics, print_trajectory

world = WorldModel(min_lane=1, max_lane=2, seg_to_m=2.0, step_duration=1.0)

planned = {
    "car1": CarState(seg=55, lane=1, prev_step=11, steps_in_lane=4, vehicle_length_segs=3),
}
oracles = {
    "ego": OracleSpec(
        init=CarState(seg=50, lane=2, prev_step=11, steps_in_lane=4, vehicle_length_segs=3),
        step_fn=make_constant_oracle(),
    ),
}
goals = {
    "car1": [
        Goal(description="end up 12–40 m behind ego in lane 2",
             ref_agent="ego", lane_eq=2, diff_min_segs=-20, diff_max_segs=-6),
    ],
}

scenario = ScenarioDefinition(
    planned=planned, oracles=oracles, goals=goals,
    world=world, constraints=ConstraintSet(), label="my scenario",
)
raw = run_mpc_simulation(scenario, SolverParams(horizon=25))

dims = {"car1": VehicleDims(4.8, 1.9), "ego": VehicleDims(4.8, 1.9)}
result = build_sim_result(raw, scenario, dims,
                          agent_colors={"car1": "#00d2ff", "ego": "#ff6b35"},
                          agent_labels={"car1": "C1", "ego": "EGO"})
print_trajectory(result)
plot_metrics(result)
animate(result)

---

IR Reference

All solver inputs are plain Python dataclasses. No file parsing required.

WorldModel

WorldModel(
    min_seg=-20, max_seg=600,   # longitudinal bounds [segs]
    min_lane=1,  max_lane=3,    # lane indices (inclusive)
    seg_to_m=2.0,               # metres per segment
    step_duration=1.0,          # seconds per solver step
)

CarState

CarState(
    seg=50,                    # longitudinal position [segs]
    lane=2,                    # lane index
    prev_step=7,               # speed at last step [segs/step]
    steps_in_lane=4,           # steps spent in current lane (used for lane change timing)
    direction=1,               # 1 = forward, -1 = oncoming
    vehicle_length_segs=3,     # vehicle length [segs]
)
# unit conversions: seg = metres / seg_to_m,  step = mps / seg_to_m

OracleSpec

OracleSpec(
    init=CarState(...),
    step_fn=make_constant_oracle(),   # what actually happens in simulation
    predict_fn=None,                  # what the planner assumes (None = constant velocity)
)

step_fn and predict_fn are deliberately separate: you can inject a real simulator as step_fn while the planner still uses a simple prediction model, or provide a smarter predict_fn (IDM, learned model) without changing the oracle.

Note on ego: In all included scenarios, ego is a dumb oracle — constant velocity or bounded-random speed, fixed lane. This is a placeholder. The intended use case is to replace ego with a real system under test (SUT): connect your AV stack's output as step_fn and the planned agents will react to whatever ego actually does.

Built-in oracle factories:

Factory	Behavior
make_constant_oracle()	constant velocity, fixed lane
make_oncoming_oracle()	constant velocity, opposite direction
make_random_oracle(min_step, max_step, delta, seed)	bounded random speed

Goal

Goal(
    description="overtake — 2–20 m ahead in lane 3",
    ref_agent="ego",          # reference agent for relative position
    lane_eq=3,                # required lane (None = any)
    diff_min_segs=1,          # min relative position [segs] (None = unbounded)
    diff_max_segs=10,         # max relative position [segs] (None = unbounded)
    permanent=False,          # False = terminal (satisfied once), True = hold forever
)

Goals are evaluated sequentially: goal i is only attempted once i−1 is done. A permanent=True goal is enforced as a hard constraint at every horizon step for the rest of the simulation.

ConstraintSet

ConstraintSet(
    kinematic=KinematicConstraints(max_step=12, min_step=0, max_step_delta=1),
    safety=SafetyConstraints(gap_segs=3, lane_change_gap_segs=4),
    lane_change=LaneChangeConstraints(change_dur=4),
)

ScenarioDefinition

ScenarioDefinition(
    planned={...},      # dict[name, CarState]   — agents controlled by CP-SAT
    oracles={...},      # dict[name, OracleSpec] — observed agents
    goals={...},        # dict[name, list[Goal]]
    world=WorldModel(...),
    constraints=ConstraintSet(...),
    label="my scenario",
    max_steps=60,
)

---

Extension Points

Custom oracle — inject a simulator or a smarter predictor:

def my_step_fn(sim_step: int, state: CarState) -> CarState:
    ...  # call external sim, return next CarState

OracleSpec(init=..., step_fn=my_step_fn, predict_fn=my_predict_fn)

Custom trajectory fitter — replace the Frenet spline with your own planner:

def my_fitter(
    segs: list[int],       # longitudinal waypoints from CP-SAT
    lanes: list[int],
    ego_lane: int,
    seg_to_m: float,
    step_duration: float,
) -> Trajectory:           # see trajectory.py
    ...

result = build_sim_result(raw, scenario, dims, fitter=my_fitter)

The fitter receives the discrete CP-SAT waypoints and is responsible for producing a continuous Trajectory. It has no effect on the solver itself.

Custom constraints — add a pure builder function following the pattern in solver.py:

def apply_my_constraints(model, name, v, ...) -> None:
    model.add(...)  # only adds to model, returns nothing

# call it inside mpc_step() after the existing builders

---

Included Scenarios

make elk       # or: uv run python -m cpsat_scenario_planner.scenarios.elk_overtaking
make merge     # or: uv run python -m cpsat_scenario_planner.scenarios.highway_merge
make overtake  # or: uv run python -m cpsat_scenario_planner.scenarios.joint_overtake

Scenario	Agents	Description
elk_overtaking	1 planned, 1 oracle	DUT brakes and changes lane; ends 12–40 m behind ego
highway_merge	2 planned, 1 oracle	Two vehicles merge from on-ramp ahead of ego
joint_overtake	2 planned, 3 oracles	car1 overtakes left, truck overtakes right, oncoming traffic

Each scenario file is self-contained: _SCENARIO_DEF is the full setup, run_simulation() runs it.

---

Running Tests

make test   # or: uv run pytest tests/ -v

---

Dependencies

• or-tools — CP-SAT solver
• numpy — trajectory spline math
• matplotlib — visualization