[skip ci] working on incorporating trajectory game example

examples/TrajectoryExamples.jl

@@ -0,0 +1,83 @@
+"""
+Utilities for constructing trajectory games, in which each player wishes to
+solve a problem of the form:
+                min_{τᵢ}   fᵢ(τ, θ)
+
+where all vehicles must jointly satisfy the constraints
+                           g̃(τ, θ) = 0
+                           h̃(τ, θ) ≥ 0.
+
+Here, τᵢ is the ith vehicle's trajectory, consisting of states and controls.
+The shared constraints g̃ and h̃ incorporate dynamic feasibility, fixed initial
+condition, actuator and state limits, environment boundaries, and
+collision-avoidance.
+"""
+
+module TrajectoryExamples
+using LevelTwoInverseGames:
+    ParametricGame, OptimizationProblem, num_players, solve
+
+using LazySets: LazySets
+using TrajectoryGamesBase:
+    TrajectoryGamesBase,
+    PolygonEnvironment,
+    ProductDynamics,
+    TimeSeparableTrajectoryGameCost,
+    TrajectoryGame,
+    GeneralSumCostStructure,
+    num_players,
+    time_invariant_linear_dynamics,
+    unstack_trajectory,
+    stack_trajectories,
+    state_dim,
+    control_dim,
+    state_bounds,
+    control_bounds,
+    OpenLoopStrategy,
+    JointStrategy,
+    RecedingHorizonStrategy,
+    rollout
+using TrajectoryGamesExamples: planar_double_integrator, animate_sim_steps
+using BlockArrays: mortar, blocks, BlockArray, Block
+using GLMakie: GLMakie
+using Makie: Makie
+using PATHSolver: PATHSolver
+using LinearAlgebra: norm_sqr, norm
+using ProgressMeter: ProgressMeter
+
+include("utils.jl")
+export build_parametric_game,
+    WarmStartRecedingHorizonStrategy,
+    pack_observations,
+    unpack_observations,
+    pack_parameters,
+    unpack_trajectory,
+    pack_parameters,
+    unpack_parameters,
+    parameter_mask,
+    generate_initial_guess
+
+"Zygote-friendly trajectory unpacking utility."
+function unpack_trajectory_zygote(flat_trajectories; dynamics::ProductDynamics)
+    horizon = Int(
+        length(flat_trajectories[1]) /
+        (state_dim(dynamics, 1) + control_dim(dynamics, 1)),
+    )
+
+    map(1:num_players(dynamics), flat_trajectories) do ii, traj
+        num_states = state_dim(dynamics, ii) * horizon
+        X = reshape(traj[1:num_states], (state_dim(dynamics, ii), horizon))
+        U = reshape(
+            traj[(num_states + 1):end],
+            (control_dim(dynamics, ii), horizon),
+        )
+
+        (; xs = eachcol(X) |> collect, us = eachcol(U) |> collect)
+    end
+end
+
+export unpack_trajectory_zygote
+
+include("lane_change.jl")
+
+end # module TrajectoryExamples

examples/lane_change.jl

@@ -0,0 +1,146 @@
+module LaneChange
+
+using ..TrajectoryExamples
+using MCPSolver
+
+using LazySets: LazySets
+using TrajectoryGamesBase:
+    TrajectoryGamesBase,
+    PolygonEnvironment,
+    ProductDynamics,
+    TimeSeparableTrajectoryGameCost,
+    TrajectoryGame,
+    GeneralSumCostStructure,
+    num_players,
+    time_invariant_linear_dynamics,
+    unstack_trajectory,
+    stack_trajectories,
+    state_dim,
+    control_dim,
+    state_bounds,
+    control_bounds,
+    OpenLoopStrategy,
+    JointStrategy,
+    RecedingHorizonStrategy,
+    rollout
+using TrajectoryGamesExamples: planar_double_integrator, animate_sim_steps
+using BlockArrays: mortar, blocks, BlockArray, Block
+using GLMakie: GLMakie
+using Makie: Makie
+using PATHSolver: PATHSolver
+using LinearAlgebra: norm_sqr, norm
+using ProgressMeter: ProgressMeter
+
+"Utility to create the road environment."
+function setup_road_environment(; lane_width = 2, num_lanes = 2, height = 50)
+    lane_centers = map(lane_idx -> (lane_idx - 0.5) * lane_width, 1:num_lanes)
+    vertices = [
+        [first(lane_centers) - 0.5lane_width, 0],
+        [last(lane_centers) + 0.5lane_width, 0],
+        [last(lane_centers) + 0.5lane_width, height],
+        [first(lane_centers) - 0.5lane_width, height],
+    ]
+
+    (; lane_centers, environment = PolygonEnvironment(vertices))
+end
+
+"Utility to set up a (two player) trajectory game."
+function setup_trajectory_game(; environment)
+    cost = let
+        stage_costs = map(1:2) do ii
+            (x, u, t, θi) -> let
+                lane_preference = last(θi)
+
+                (x[Block(ii)][1] - lane_preference)^2 +
+                0.5norm_sqr(x[Block(ii)][3:4] - [0, 2]) +
+                0.1norm_sqr(u[Block(ii)])
+            end
+        end
+
+        function reducer(stage_costs)
+            reduce(+, stage_costs) / length(stage_costs)
+        end
+
+        TimeSeparableTrajectoryGameCost(
+            stage_costs,
+            reducer,
+            GeneralSumCostStructure(),
+            1.0,
+        )
+    end
+
+    function coupling_constraints(xs, us, θ)
+        mapreduce(vcat, xs) do x
+            x1, x2 = blocks(x)
+
+            # Players need to stay at least 2 m away from one another.
+            norm_sqr(x1[1:2] - x2[1:2]) - 4
+        end
+    end
+
+    agent_dynamics = planar_double_integrator(;
+        state_bounds = (; lb = [-Inf, -Inf, -10, 0], ub = [Inf, Inf, 10, 10]),
+        control_bounds = (; lb = [-5, -5], ub = [3, 3]),
+    )
+    dynamics = ProductDynamics([agent_dynamics for _ in 1:2])
+
+    TrajectoryGame(dynamics, cost, environment, coupling_constraints)
+end
+
+using Infiltrator
+
+function main(;
+    initial_state = mortar([[1.0, 1.0, 0.0, 1.0], [3.2, 0.9, 0.0, 1.0]]),
+    horizon = 10,
+    height = 50.0,
+    num_lanes = 2,
+    lane_width = 2,
+    num_sim_steps = 150,
+    gradient_steps_per_turn = 5,
+    observation_noise_stddev = 0.1,
+)
+    (; environment, lane_centers) =
+        setup_road_environment(; num_lanes, lane_width, height)
+    game = LaneChange.setup_trajectory_game(; environment)
+
+    # Build a game. Each player has a parameter for lane preference.
+    # P1 wants to stay in the left lane, and P2 wants to move from the
+    # right to the left lane.
+    lane_preferences = mortar([[lane_centers[1]], [lane_centers[1]]])
+    mcp = build_mcp(; game, horizon, params_per_player = 1)
+
+    # Simulate the ground truth.
+    turn_length = 3
+    sim_steps = let
+        progress = ProgressMeter.Progress(num_sim_steps)
+        ground_truth_strategy = WarmStartRecedingHorizonStrategy(;
+            game,
+            parametric_game,
+            turn_length,
+            horizon,
+            parameters = lane_preferences,
+        )
+
+        rollout(
+            game.dynamics,
+            ground_truth_strategy,
+            initial_state,
+            num_sim_steps;
+            get_info = (γ, x, t) ->
+                (ProgressMeter.next!(progress); γ.receding_horizon_strategy),
+        )
+    end
+
+    animate_sim_steps(
+        game,
+        sim_steps;
+        live = false,
+        framerate = 20,
+        show_turn = true,
+        xlims = (first(lane_centers) - lane_width, last(lane_centers) + lane_width),
+        ylims = (-1, height + 1),
+        aspect = num_lanes * lane_width / height,
+    )
+end
+
+end # module LaneChange