From f99f19cc31fcd88b177b6abeef16dddebe0665f4 Mon Sep 17 00:00:00 2001 From: Fernando Palafox Date: Fri, 13 Oct 2023 12:54:01 -0600 Subject: [PATCH] cleaned up and built dzdq --- experiments/tower_defense.jl | 289 +++++++---------------------------- 1 file changed, 56 insertions(+), 233 deletions(-) diff --git a/experiments/tower_defense.jl b/experiments/tower_defense.jl index 5abc3f6..dd7f8fb 100644 --- a/experiments/tower_defense.jl +++ b/experiments/tower_defense.jl @@ -16,249 +16,72 @@ using Zygote """ -function solve_tower_defense_game(θ, pₖ) - - fs = [(x,θ) -> -(pₖ[1]*x[Block(1)]'*x[Block(2)] + pₖ[2]*x[Block(1)]'*x[Block(3)] + pₖ[3]*x[Block(1)]'*x[Block(4)]), - (x,θ) -> x[Block(1)]'*diagm(θ[1:3])*x[Block(2)], - (x,θ) -> x[Block(1)]'*diagm(θ[4:6])*x[Block(3)], - (x,θ) -> x[Block(1)]'*diagm(θ[7:9])*x[Block(4)]] - - function g1(x, θ) - [sum(x[Block(1)]) - 1] - end - function g2(x, θ) - [sum(x[Block(2)]) - 1] - end - function g3(x, θ) - [sum(x[Block(3)]) - 1] - end - function g4(x, θ) - [sum(x[Block(4)]) - 1] - end - gs = [g1, g2, g3, g4] - - function h1(x, θ) - x[Block(1)] - end - function h2(x, θ) - x[Block(2)] - end - function h3(x, θ) - x[Block(3)] - end - function h4(x, θ) - x[Block(4)] - end - hs = [h1, h2, h3, h4] - g̃ = (x, θ) -> [0] - h̃ = (x, θ) -> [0] - - #Main.@infiltrate - - problem = ParametricGame(; - objectives = fs, - equality_constraints = gs, - inequality_constraints = hs, - shared_equality_constraint = g̃, - shared_inequality_constraint = h̃, - parameter_dimension = 9, - primal_dimensions = [3,3,3,3], - equality_dimensions = [1,1,1,1], - inequality_dimensions = [3,3,3,3], - shared_equality_dimension = 1, - shared_inequality_dimension = 1, - ) - - (;solution = solve(problem, parameter_value = θ), fs) -end - -function solve_full_tower_defense_game(θ, pₖ) - - fs = [# Stage 2 - (x,θ) -> -(x[Block(9)][1]*pₖ[1]*x[Block(1)]'*x[Block(2)] + x[Block(9)][2]*pₖ[2]*x[Block(1)]'*x[Block(3)] + x[Block(9)][3]*pₖ[3]*x[Block(1)]'*x[Block(4)]) / (x[Block(9)]'*pₖ), - (x,θ) -> x[Block(1)]'*diagm(θ[1:3])*x[Block(2)], - (x,θ) -> x[Block(1)]'*diagm(θ[4:6])*x[Block(3)], - (x,θ) -> x[Block(1)]'*diagm(θ[7:9])*x[Block(4)], - (x,θ) -> -((1-x[Block(9)][1])*pₖ[1]*x[Block(5)]'*x[Block(6)] + (1-x[Block(9)][2])*pₖ[2]*x[Block(5)]'*x[Block(7)] + (1-x[Block(9)][3])*pₖ[3]*x[Block(5)]'*x[Block(8)]) / (1 - x[Block(9)]'*pₖ), - (x,θ) -> x[Block(5)]'*diagm(θ[1:3])*x[Block(6)], - (x,θ) -> x[Block(5)]'*diagm(θ[4:6])*x[Block(7)], - (x,θ) -> x[Block(5)]'*diagm(θ[7:9])*x[Block(8)], - - #Stage 1: Weighted sum of P1 costs, weighted by P(s¹=1) and P(s¹=0) - (x,θ) -> -(pₖ[1]*(x[Block(9)][1]*(x[Block(1)]'*x[Block(2)] - x[Block(5)]'*x[Block(6)]) + x[Block(5)]'*x[Block(6)]) \ - + pₖ[2]*(x[Block(9)][2]*(x[Block(1)]'*x[Block(3)] - x[Block(5)]'*x[Block(7)]) + x[Block(5)]'*x[Block(7)]) \ - + pₖ[3]*(x[Block(9)][3]*(x[Block(1)]'*x[Block(4)] - x[Block(5)]'*x[Block(8)]) + x[Block(5)]'*x[Block(8)])) - ] - - # Equality Constraints - function g1(x, θ) - [sum(x[Block(1)]) - 1] - end - function g2(x, θ) - [sum(x[Block(2)]) - 1] - end - function g3(x, θ) - [sum(x[Block(3)]) - 1] - end - function g4(x, θ) - [sum(x[Block(4)]) - 1] - end - function g5(x, θ) - [sum(x[Block(5)]) - 1] - end - function g6(x, θ) - [sum(x[Block(6)]) - 1] - end - function g7(x, θ) - [sum(x[Block(7)]) - 1] - end - function g8(x, θ) - [sum(x[Block(8)]) - 1] - end - function g9(x, θ) # test: qₖ = 1 ∀k or put placeholder? - [ - #x[Block(9)][1] - 1; - #x[Block(9)][2] - 1; - #x[Block(9)][3] - 1; - sum(x[Block(8)]) - 1 - ] - end - gs = [g1, g2, g3, g4, g5, g6, g7, g8, g9] - - # Inequality Constraints - function h1(x, θ) - x[Block(1)] - end - function h2(x, θ) - x[Block(2)] - end - function h3(x, θ) - x[Block(3)] - end - function h4(x, θ) - x[Block(4)] - end - function h5(x, θ) - x[Block(5)] - end - function h6(x, θ) - x[Block(6)] - end - function h7(x, θ) - x[Block(7)] - end - function h8(x, θ) - x[Block(8)] - end - function h9(x, θ) - [ - x[Block(9)]; - ones(size(x[Block(9)])) - x[Block(9)]; - #1.5 - sum(x[Block(9)]) - ] - end - hs = [h1, h2, h3, h4, h5, h6, h7, h8, h9] - g̃ = (x, θ) -> [0] - h̃ = (x, θ) -> [0] - - #Main.@infiltrate - - problem = ParametricGame(; - objectives = fs, - equality_constraints = gs, - inequality_constraints = hs, - shared_equality_constraint = g̃, - shared_inequality_constraint = h̃, - parameter_dimension = 9, # θ - primal_dimensions = [3,3,3,3,3,3,3,3,3], # x - equality_dimensions = [1,1,1,1,1,1,1,1,1], - inequality_dimensions = [3,3,3,3,3,3,3,3,6], - shared_equality_dimension = 1, - shared_inequality_dimension = 1, - ) - - (;solution = solve(problem, parameter_value = θ), fs) -end - -function solve_full_tower_defense_game_stage2(θ, pₖ, wₖ) - - p_1 = θ' * pₖ - p_0 = (1 .- θ)' * pₖ - - # objectives +function build_parametric_game(pₖ, wₖ) fs = [ - (x,θ) -> -(θ[1]*pₖ[1]*x[Block(1)]'*x[Block(2)] + θ[2]*pₖ[2]*x[Block(1)]'*x[Block(3)] + θ[3]*pₖ[3]*x[Block(1)]'*x[Block(4)])/p_1, - (x,θ) -> x[Block(1)]'*diagm(wₖ[1:3])*x[Block(2)], - (x,θ) -> x[Block(1)]'*diagm(wₖ[4:6])*x[Block(3)], - (x,θ) -> x[Block(1)]'*diagm(wₖ[7:9])*x[Block(4)], - (x,θ) -> -((1 - θ[1])*pₖ[1]*x[Block(5)]'*x[Block(6)] + (1 - θ[2])*pₖ[2]*x[Block(5)]'*x[Block(7)] + (1 - θ[3])*pₖ[3]*x[Block(5)]'*x[Block(8)])/p_0, - (x,θ) -> x[Block(5)]'*diagm(wₖ[1:3])*x[Block(6)], - (x,θ) -> x[Block(5)]'*diagm(wₖ[4:6])*x[Block(7)], - (x,θ) -> x[Block(5)]'*diagm(wₖ[7:9])*x[Block(8)], - ] + (x, θ) -> -(θ[1] * pₖ[1] * x[Block(1)]' * x[Block(2)] + θ[2] * pₖ[2] * x[Block(1)]' * x[Block(3)] + θ[3] * pₖ[3] * x[Block(1)]' * x[Block(4)]) / (θ' * pₖ), + (x, θ) -> x[Block(1)]' * diagm(wₖ[1:3]) * x[Block(2)], + (x, θ) -> x[Block(1)]' * diagm(wₖ[4:6]) * x[Block(3)], + (x, θ) -> x[Block(1)]' * diagm(wₖ[7:9]) * x[Block(4)], + (x, θ) -> -((1 - θ[1]) * pₖ[1] * x[Block(5)]' * x[Block(6)] + (1 - θ[2]) * pₖ[2] * x[Block(5)]' * x[Block(7)] + (1 - θ[3]) * pₖ[3] * x[Block(5)]' * x[Block(8)]) / ((1 .- θ)' * pₖ), + (x, θ) -> x[Block(5)]' * diagm(wₖ[1:3]) * x[Block(6)], + (x, θ) -> x[Block(5)]' * diagm(wₖ[4:6]) * x[Block(7)], + (x, θ) -> x[Block(5)]' * diagm(wₖ[7:9]) * x[Block(8)], + ] # equality constraints - gs = [(x,θ) -> sum(x[Block(i)]) - 1 for i in 1:8] + gs = [(x, θ) -> [sum(x[Block(i)]) - 1] for i in 1:8] # inequality constraints - hs = [(x,θ) -> x[Block(i)] for i in 1:8] + hs = [(x, θ) -> x[Block(i)] for i in 1:8] # shared constraints g̃ = (x, θ) -> [0] h̃ = (x, θ) -> [0] - - problem = ParametricGame(; - objectives = fs, - equality_constraints = gs, - inequality_constraints = hs, - shared_equality_constraint = g̃, - shared_inequality_constraint = h̃, - parameter_dimension = 3, - primal_dimensions = [3,3,3,3,3,3,3,3], - equality_dimensions = [1,1,1,1,1,1,1,1], - inequality_dimensions = [3,3,3,3,3,3,3,3], - shared_equality_dimension = 1, - shared_inequality_dimension = 1, - ) - - (;solution = solve(problem, parameter_value = θ), fs) -end - -# VoI -function VoI_TD() - - # Worlds - wₖ = [0, 1, 1, 1, 0, 1, 1, 1, 0] # P2 cares about a single direction in each world - # Prior distribution of k worlds - pₖ = [1/3; 1/3; 1/3] - - # Stage 1 - q = [1/3; 1/3; 1/3] # replace w/ solution from stage 1 - - # Stage 2 - results = solve_full_tower_defense_game_stage2(q, pₖ, wₖ) - vars = BlockArray(results.solution.z[1:24], [3,3,3,3,3,3,3,3]) - - # TODO Stage 1 - - # Stage 2 - P1_cost_stage2_s1 = results.fs[1](vars, q) - P1_cost_stage2_s0 = results.fs[5](vars, q) - - # println("P1 Cost of tower defense game (stage 1): ", P1_cost_stage1) - println("P1 Cost of tower defense game (stage 2, signal 1): ", P1_cost_stage2_s1) - println("P1 Cost of tower defense game (stage 2, signal 0): ", P1_cost_stage2_s0) - - # TEMPORARY: - # Zygote.jacobian(q -> solve_full_tower_defense_game_stage2(q, pₖ, wₖ).solution.z[1:24], q) + ParametricGame(; + objectives=fs, + equality_constraints=gs, + inequality_constraints=hs, + shared_equality_constraint=g̃, + shared_inequality_constraint=h̃, + parameter_dimension=3, + primal_dimensions=[3, 3, 3, 3, 3, 3, 3, 3], + equality_dimensions=[1, 1, 1, 1, 1, 1, 1, 1], + inequality_dimensions=[3, 3, 3, 3, 3, 3, 3, 3], + shared_equality_dimension=1, + shared_inequality_dimension=1 + ), fs end -"Jacobian of Stage 2 solution w.r.t. P1's signal structure" -function stage_2_Jacobian(q, pₖ, wₖ) - # Solve Stage 2 - results = solve_full_tower_defense_game_stage2(q, pₖ, wₖ) - - # Compute Jacobian of solution - +function_ + +"""Solve Stackelberg-like game with 2 stages. + + Returns dz/dq: Jacobian of Stage 2's decision variable (z = P1 and P2's variables in a Bayesian game) w.r.t. Stage 1's decision variable (q = signal structure)""" +function dzdq() + # Setup game + wₖ = [0, 1, 1, 1, 0, 1, 1, 1, 0] # worlds + pₖ = [1/3; 1/3; 1/3] # P1's prior distribution over worlds + parametric_game, fs = build_parametric_game(pₖ, wₖ) + + # Solve Stage 1 + q = [1/3; 1/3; 1/3] # TODO obtain q by solving Stage 1 + + # Solve Stage 2 given q + solution = solve( + parametric_game, + q; + initial_guess = zeros(total_dim(parametric_game)), + verbose=true, + ) + vars = BlockArray(solution.variables[1:24], [3,3,3,3,3,3,3,3]) + + # Return Jacobian + Zygote.jacobian(q -> solve( + parametric_game, + q; + initial_guess=zeros(total_dim(parametric_game)), + verbose=false, + return_primals=false + ).variables[1:24], q)[1] end \ No newline at end of file