diff --git a/julia_code/.vscode/extensions.json b/julia_code/.vscode/extensions.json new file mode 100644 index 0000000..d1987c7 --- /dev/null +++ b/julia_code/.vscode/extensions.json @@ -0,0 +1,13 @@ +{ + // See https://go.microsoft.com/fwlink/?LinkId=827846 to learn about workspace recommendations. + // Extension identifier format: ${publisher}.${name}. Example: vscode.csharp + + // List of extensions which should be recommended for users of this workspace. + "recommendations": [ + + ], + // List of extensions recommended by VS Code that should not be recommended for users of this workspace. + "unwantedRecommendations": [ + + ] +} \ No newline at end of file diff --git a/julia_code/.vscode/settings.json b/julia_code/.vscode/settings.json new file mode 100644 index 0000000..25c729f --- /dev/null +++ b/julia_code/.vscode/settings.json @@ -0,0 +1,3 @@ +{ + "python.pythonPath": "/bin/python3" +} \ No newline at end of file diff --git a/julia_code/Module/notebooks/ExampleUse.jl b/julia_code/Module/notebooks/ExampleUse.jl new file mode 100644 index 0000000..40ebaac --- /dev/null +++ b/julia_code/Module/notebooks/ExampleUse.jl @@ -0,0 +1,81 @@ +include("../src/Orbits.jl") +using .Orbits +using LinearAlgebra + +#Set key dimensions +const N_constellations = 4 +const N_debris = 1 +const N_states = N_constellations + N_debris + +#Setup Economic Models +em2_a = LinearModel(0.95, [1 -0.02 -0.02 0], [5.0 0 0 0]) +em2_b = LinearModel(0.95, [-0.02 1 -0.02 0], [0.0 5 0 0]) +em2_c = LinearModel(0.95, [0 -0.02 1 -0.02], [0.0 0 5 0]) +em2_d = LinearModel(0.95, [0 -0.02 -0.02 1], [0.0 0 0 5]) + +#Setup Physics + +basic_model = BasicPhysics( + 0.002 + ,0.002*(ones(N_constellations,N_constellations) - LinearAlgebra.I) + ,0.01 + ,0.001 + ,5.0 + ,0.05 +) + +# Setup NN +bg = BranchGenerator(N_constellations) +operators_policy = bg(operator_policy_function_generator(N_constellations,N_debris),vcat) +planners_policy = bg(operator_policy_function_generator(N_constellations,N_debris),vcat) +planners_value = value_function_generator(64) + + +# Setup Operators +const operator_array = [ + #first operator + OperatorLoss( + em2_a + ,value_function_generator() + ,operators_policy #this is held by all operators + ,params(operators_policy[2][1]) #first operator gets first branch of params + ,basic_model + ) + ,OperatorLoss( + em2_b + ,value_function_generator() + ,operators_policy #this is held by all operators + ,params(operators_policy[2][2]) #first operator gets first branch of params + ,basic_model + ) + ,OperatorLoss( + em2_c + ,value_function_generator() + ,operators_policy #this is held by all operators + ,params(operators_policy[2][3]) #first operator gets first branch of params + ,basic_model + ) + ,OperatorLoss( + em2_d + ,value_function_generator() + ,operators_policy #this is held by all operators + ,params(operators_policy[2][4]) #first operator gets first branch of params + ,basic_model + ) +] + +#sanity check time +@assert length(operator_array) == N_constellations "Mismatch in predetermined number of constellations and the number of operators initialized" + +# Setup Planner +pl = PlannerLoss( + 0.95 + ,operator_array + ,planners_policy + ,params(planners_policy) + ,planners_value + ,params(planners_value) + ,basic_model +) + +# Export Planner \ No newline at end of file diff --git a/julia_code/Module/src/Constellations.jl b/julia_code/Module/src/Constellations.jl new file mode 100644 index 0000000..476f621 --- /dev/null +++ b/julia_code/Module/src/Constellations.jl @@ -0,0 +1,401 @@ + +module Orbits + + +#Add exports here +export State ,state_to_tuple ,state_transition ,PhysicalModel ,BasicPhysics + +export BranchGenerator, value_function_generator, + cross_linked_planner_policy_function_generator, + operator_policy_function_generator + +# Exports +export GeneralizedLoss ,OperatorLoss ,PlannerLoss ,UniformDataConstructor, + LinearModel ,BasicPhysics, + operator_policy_function_generator ,value_function_generator ,BranchGenerator, + TupleDuplicator, BranchGenerator + +using Flux, LinearAlgebra +using BSON, Zygote +export LinearModel, BenefitFunction, EconomicModel + +#==# +struct State + stocks::Array{Float32} + debris::Array{Float32} +end + +function state_to_tuple(s::State) + return (s.stocks ,s.debris) +end + + + +abstract type PhysicalModel end + +struct BasicPhysics <: PhysicalModel + #rate at which debris hits satellites + debris_collision_rate::Real + #rate at which satellites of different constellations collide + satellite_collision_rates::Matrix{Float64} + #rate at which debris exits orbits + decay_rate::Real + #rate at which satellites + autocatalysis_rate::Real + #ratio at which a collision between satellites produced debris + satellite_collision_debris_ratio::Real + #Ratio at which launches produce debris + launch_debris_ratio::Real +end +function state_transition( + physics::BasicPhysics + ,state::State + ,launches::Vector{Float32} + ) + #= + Physical Transitions + =# + survival_rate = survival(state,physics) + + # Debris transitions + + # get changes in debris from natural dynamics + natural_debris_dynamics = (1 - physics.decay_rate + physics.autocatalysis_rate) * state.debris + + # get changes in debris from satellite loss + satellite_loss_debris = physics.satellite_collision_debris_ratio * (1 .- survival_rate)' * state.stocks + + # get changes in debris from launches + launch_debris = physics.launch_debris_ratio * sum(launches) + + # total debris level + debris′ = natural_debris_dynamics .+ satellite_loss_debris .+ launch_debris + + + # Stocks Transitions + stocks′ = (LinearAlgebra.diagm(survival_rate) .- physics.decay_rate)*state.stocks + launches + + + return State(stocks′,debris′) +end + + +function survival( + state::State + ,physical_model::BasicPhysics + ) + return exp.(-(physical_model.satellite_collision_rates .+ physical_model.decay_rate) * state.stocks .- (physical_model.debris_collision_rate * state.debris)) +end + + + + +#= +Benefit Functions: + +These represent the benefits that different operators may recieve +=# +abstract type EconomicModel end + +#basic linear model +struct LinearModel <: EconomicModel + β::Float32 + payoff_array::Array{Float32} + policy_costs::Array{Float32} +end +function (em::LinearModel)( + state::State + ,actions::Vector{Float32} + ) + return em.payoff_array*state.stocks - em.policy_costs*actions +end + +#basic CES model +struct CES <: EconomicModel + β::Float32 + r::Float32 #elasticity of subsititution + payoff_array::Array{Float32} + policy_costs::Array{Float32} + debris_costs::Array{Float32} +end +function (em::CES)( + state::State + ,actions::Vector{Float32} + ) + + #issue here with multiplication + r1 = em.payoff_array .* (state.stocks.^em.r) + r2 = - em.debris_costs .* (state.debris.^em.r) + r3 = - em.policy_costs .* (actions.^em.r) + return (r1 + r2 + r3) .^ (1/em.r) + +end + +#basic CRRA +struct CRRA <: EconomicModel + β::Float32 + σ::Float32 #elasticity of subsititution + payoff_array::Array{Float32} + policy_costs::Array{Float32} + debris_costs::Array{Float32} +end +function (em::CRRA)( + state::State + ,actions::Vector{Float32} + ) + #issue here with multiplication + core = (em.payoff_array*state.stocks - em.debris_costs*state.debris - em.policy_costs*actions).^(1 - em.σ) + + return (core-1)/(1-em.σ) +end + + + +#= TupleDuplicator +This is used to create a tuple of size n with deepcopies of any object x +=# +struct TupleDuplicator + n::Int +end +#make tuples consisting of copies of whatever was provided +function (f::TupleDuplicator)(s::State) + st = state_to_tuple(s) + return f(st) + #BROKEN: Fails in test, but works in REPL. +end + +(f::TupleDuplicator)(x) = tuple([deepcopy(x) for i=1:f.n]...) + + + +struct BranchGenerator + n::UInt +end +function (b::BranchGenerator)(branch::Flux.Parallel,join_fn::Function) + # used to deepcopy the branches and duplicate the inputs in the returned chain + f = TupleDuplicator(b.n) + + return Flux.Chain(state_to_tuple,f,Flux.Parallel(join_fn,f(branch)...)) + #note that it destructures the state to a tuple, duplicates, + # and then passes to the parallelized functions. +end + + +# Neural Network Generators + +function value_function_generator(number_params=32) + return Flux.Chain( + Flux.Parallel(vcat + #parallel joins together stocks and debris, after a little bit of preprocessing + ,Flux.Chain( + Flux.Dense(N_constellations, number_params*2,Flux.relu) + ,Flux.Dense(number_params*2, number_params*2,Flux.σ) + ) + ,Flux.Chain( + Flux.Dense(N_debris, number_params,Flux.relu) + ,Flux.Dense(number_params, number_params,Flux.σ) + ) + ) + #Apply some transformations to the preprocessed data. + ,Flux.Dense(number_params*3,number_params,Flux.σ) + ,Flux.Dense(number_params,number_params,Flux.σ) + ,Flux.Dense(number_params,1) + ) + +end + +function cross_linked_planner_policy_function_generator(number_params=32) + return Flux.Chain( + Flux.Parallel(vcat + #parallel joins together stocks and debris + ,Flux.Chain( + Flux.Dense(N_constellations, number_params*2,Flux.relu) + ,Flux.Dense(number_params*2, number_params*2,Flux.σ) + ) + ,Flux.Chain( + Flux.Dense(N_debris, number_params,Flux.relu) + ,Flux.Dense(number_params, number_params) + ) + ) + #Apply some transformations + ,Flux.Dense(number_params*3,number_params,Flux.σ) + ,Flux.Dense(number_params,number_params,Flux.σ) + ,Flux.Dense(number_params,N_constellations,Flux.relu) + ) + +end + +function operator_policy_function_generator( + N_constellations::Int + ,N_debris + ,number_params=32 + ) + return Flux.Chain( + Flux.Parallel(vcat + #parallel joins together stocks and debris + ,Flux.Chain( + Flux.Dense(N_constellations, number_params*2,Flux.relu) + ,Flux.Dense(number_params*2, number_params*2,Flux.tanh) + ) + ,Flux.Chain( + Flux.Dense(N_debris, number_params,Flux.relu) + ,Flux.Dense(number_params, number_params) + ) + ) + #Apply some transformations + ,Flux.Dense(number_params*3,number_params,Flux.tanh) + ,Flux.Dense(number_params,number_params,Flux.tanh) + ,Flux.Dense(number_params,1,x -> Flux.relu.(sinh.(x))) + ) + +end + + + + + +abstract type GeneralizedLoss end + +#= +This struct organizes information about a given constellation operator +Used to provide an interable loss function for training +=# +struct OperatorLoss <: GeneralizedLoss + #econ model describing operator + econ_model::EconomicModel + #Operator's value and policy functions, as well as which parameters the operator can train. + operator_value_fn::Flux.Chain + collected_policies::Flux.Chain #this is held by all operators + operator_policy_params::Flux.Params + physics::PhysicalModel +end +# overriding function to calculate operator loss +function (operator::OperatorLoss)( + state::State +) + + + #get actions + a = operator.collected_policies((s,d)) + + #get updated stocks and debris + state′ = state_transition(operator.physics,state,a) + + + bellman_residuals = operator.operator_value_fn(state) - operator.econ_model(state,a) - operator.econ_model.β * operator.operator_value_fn(state′) + + maximization_condition = - operator.econ_model(state ,a) - operator.econ_model.β * operator.operator_value_fn((state′)) + + return Flux.mae(bellman_residuals.^2 ,maximization_condition) +end # struct +function (operator::OperatorLoss)( + state::State + ,policy::Flux.Chain #allow for another policy to be subsituted + ) + + + #get actions + a = policy(state) + + #get updated stocks and debris + state′ = stake_transition(state,a) + + + bellman_residuals = operator.operator_value_fn(state) - operator.econ_model(state,a) - operator.econ_model.β * operator.operator_value_fn(state′) + + maximization_condition = - operator.econ_model(state,a) - operator.econ_model.β * operator.operator_value_fn(state′) + + return Flux.mae(bellman_residuals.^2 ,maximization_condition) +end #function + +function train_operators!(op::OperatorLoss, data::State, opt) + Flux.train!(op, op.operator_policy_params, data, opt) + Flux.train!(op, Flux.params(op.operator_value_fn), data, opt) +end + + +#= +Describe the Planners loss function +=# +struct PlannerLoss <: GeneralizedLoss + #= + Ideally, with just a well formed PlannerLoss and the training functions below, we should be able to train the approximation. + + There is an issue with appropriately training the value functions. + In this case, it is not happening... + =# + β::Float32 + + operators::Array{GeneralizedLoss} + + policy::Flux.Chain + policy_params::Flux.Params + value::Flux.Chain + value_params::Flux.Params + + physical_model::PhysicalModel +end +function (planner::PlannerLoss)( + state::State +) + a = planner.policy((s ,d)) + #get updated stocks and debris + state′ = state_transition(s ,d ,a) + + + + #calculate the total benefit from each of the models + benefit = sum([ co.econ_model(s ,d ,a) for co in planner.operators]) + #issue here with mutating. Maybe generators/list comprehensions? + + + bellman_residuals = planner.value((s,d)) - benefit - planner.β .* planner.value((s′,d′)) + + maximization_condition = - benefit - planner.β .* planner.value((s′,d′)) + + return Flux.mae(bellman_residuals.^2 ,maximization_condition) +end + +function train_planner!(pl::PlannerLoss, N_epoch::Int, opt) + errors = [] + for i = 1:N_epoch + data = data_gen() + Flux.train!(pl, pl.policy_params, data, opt) + Flux.train!(pl, pl.value_params, data, opt) + append!(errors, error(pl, data1) / 200) + end + return errors +end + + +function train_operators!(pl::PlannerLoss, N_epoch::Int, opt) + errors = [] + for i = 1:N_epoch + data = data_gen() + for op in pl.operators + train_operators!(op,data,opt) + end + end + return errors +end + + +#Construct and manage data +abstract type DataConstructor end + +struct UniformDataConstructor <: DataConstructor + N::UInt64 + satellites_bottom::Float32 + satellites_top::Float32 + debris_bottom::Float32 + debris_top::Float32 +end +function (dc::UniformDataConstructor)() + #currently ignores the quantity of data it should construct. + State( + rand(dc.satellites_bottom:dc.satellites_top, N_constellations) + , rand(dc.debris_bottom:dc.debris_top, N_debris) + ) +end + +end # end module diff --git a/julia_code/Module/src/Orbits.jl b/julia_code/Module/src/Orbits.jl index 385d3a4..a165c1a 100644 --- a/julia_code/Module/src/Orbits.jl +++ b/julia_code/Module/src/Orbits.jl @@ -1,19 +1,22 @@ module Orbits -using Flux,LinearAlgebra,Zygote +using Flux,LinearAlgebra,Zygote,BSON include("physical_model.jl") -using .PhysicsModule +export State, state_to_tuple, state_transition, BasicPhysics, PhysicalModel include("benefit_functions.jl") -using .BenefitFunctions +export LinearModel, EconomicModel include("flux_helpers.jl") -using .NNTools - -# Exports -export GeneralizedLoss , OperatorLoss, PlannerLoss, UniformDataConstructor - +export operator_policy_function_generator, + value_function_generator, + BranchGenerator, + cross_linked_planner_policy_function_generator + +# Exports from below +export GeneralizedLoss ,OperatorLoss ,PlannerLoss ,UniformDataConstructor, + train_planner!, train_operators! # Code @@ -160,4 +163,5 @@ function (dc::UniformDataConstructor)() , rand(dc.debris_bottom:dc.debris_top, N_debris) ) end + end # end module diff --git a/julia_code/Module/src/benefit_functions.jl b/julia_code/Module/src/benefit_functions.jl index 1957891..8317f8c 100644 --- a/julia_code/Module/src/benefit_functions.jl +++ b/julia_code/Module/src/benefit_functions.jl @@ -1,11 +1,5 @@ +#don't write as a module -module BenefitFunctions - -export LinearModel, BenefitFunction, EconomicModel - - -include("physical_model.jl") -using .PhysicsModule: State #= Benefit Functions: @@ -66,4 +60,3 @@ function (em::CRRA)( return (core-1)/(1-em.σ) end -end #end module diff --git a/julia_code/Module/src/flux_helpers.jl b/julia_code/Module/src/flux_helpers.jl index 4503b88..0ebe699 100644 --- a/julia_code/Module/src/flux_helpers.jl +++ b/julia_code/Module/src/flux_helpers.jl @@ -1,12 +1,3 @@ -module NNTools - -export BranchGenerator, value_function_generator - -using Flux, BSON, Zygote - -include("physical_model.jl") -using .PhysicsModule: State,tuple - #= TupleDuplicator @@ -16,25 +7,26 @@ struct TupleDuplicator n::Int end #make tuples consisting of copies of whatever was provided +function (f::TupleDuplicator)(s::State) + st = state_to_tuple(s) + return f(st) + #BROKEN: Fails in test, but works in REPL. +end + (f::TupleDuplicator)(x) = tuple([deepcopy(x) for i=1:f.n]...) -#do the same but for the case of states, coerce them to tuples first -(f::TupleDuplicator)(x::State) = tuple([deepcopy(tuple(x)) for i=1:f.n]...) -#= BranchGenerator -This generates a policy function full of branches with the properly scaled sides -=# struct BranchGenerator - n::UInt8 - #limit to 2^8 operators + n::UInt end -function (b::BranchGenerator)(branch::Flux.Chain,join_fn::Function) +function (b::BranchGenerator)(branch::Flux.Parallel,join_fn::Function) # used to deepcopy the branches and duplicate the inputs in the returned chain f = TupleDuplicator(b.n) - return Flux.Chain(f,Flux.Parallel(join_fn,f(branch))) - #note that it both uses and runs the tuple duplicator f + return Flux.Chain(state_to_tuple,f,Flux.Parallel(join_fn,f(branch)...)) + #note that it destructures the state to a tuple, duplicates, + # and then passes to the parallelized functions. end @@ -82,7 +74,11 @@ function cross_linked_planner_policy_function_generator(number_params=32) end -function operator_policy_function_generator(number_params=32) +function operator_policy_function_generator( + N_constellations::Int + ,N_debris + ,number_params=32 + ) return Flux.Chain( Flux.Parallel(vcat #parallel joins together stocks and debris @@ -103,4 +99,3 @@ function operator_policy_function_generator(number_params=32) end -end diff --git a/julia_code/Module/src/physical_model.jl b/julia_code/Module/src/physical_model.jl index 2e0f8e5..c6896bc 100644 --- a/julia_code/Module/src/physical_model.jl +++ b/julia_code/Module/src/physical_model.jl @@ -1,20 +1,14 @@ -module PhysicsModule -#Add exports here -export State ,tuple ,state_transition ,PhysicalModel ,BasicModel -using Flux, LinearAlgebra -using Zygote #included for saving and loading models - #==# struct State stocks::Array{Float32} debris::Array{Float32} end -function tuple(st::State) - return (st.stocks, st.debris) +function state_to_tuple(s::State) + return (s.stocks ,s.debris) end ### Physics @@ -22,7 +16,7 @@ end abstract type PhysicalModel end -struct BasicModel <: PhysicalModel +struct BasicPhysics <: PhysicalModel #rate at which debris hits satellites debris_collision_rate::Real #rate at which satellites of different constellations collide @@ -37,7 +31,7 @@ struct BasicModel <: PhysicalModel launch_debris_ratio::Real end function state_transition( - physics::BasicModel + physics::BasicPhysics ,state::State ,launches::Vector{Float32} ) @@ -71,9 +65,8 @@ end function survival( state::State - ,physical_model::BasicModel + ,physical_model::BasicPhysics ) return exp.(-(physical_model.satellite_collision_rates .+ physical_model.decay_rate) * state.stocks .- (physical_model.debris_collision_rate * state.debris)) end -end #end module diff --git a/julia_code/Module/tests/test_NNTools.jl b/julia_code/Module/tests/test_NNTools.jl new file mode 100644 index 0000000..a8f2961 --- /dev/null +++ b/julia_code/Module/tests/test_NNTools.jl @@ -0,0 +1,54 @@ +using Test,Flux + +include("../src/Orbits.jl") +using .Orbits + + + +#= +Structure: + +This is broken into three parts + - The test_interfaces module contains various tools used in testing + - The test_routines includes functions that setup, run, and teardown tests + - Everything else is a set of tests using the standard testing tools in Julia +=# +@testset "Overall" verbose=true begin +@testset "TupleDuplicator" verbose=true begin + #Check if tuple duplicator duplicates something + td2 = Orbits.TupleDuplicator(2) + @test typeof(td2) <: Orbits.TupleDuplicator + + @test td2(([1,2],[3])) == (([1,2],[3]),([1,2],[3])) + + st = State([1.0,2],[3]) + @test td2((st.stocks,st.debris)) == (([1.0,2],[3]),([1.0,2],[3])) + + @test td2(state_to_tuple(st)) == (([1.0,2],[3]),([1.0,2],[3])) + + @test td2(st) == ((st.stocks,st.debris),(st.stocks,st.debris)) +end + +@testset "BranchGenerator" verbose=true begin + st = State([1.0,2],[3]) + tp = state_to_tuple(st) + bg = BranchGenerator(2) + + branch = Flux.Parallel(vcat, Dense(2,1),Dense(1,1)) + + branched = bg(branch,vcat) + + #type tests + @test typeof(branch(tp)) <: Array{Float32} + + @test_broken typeof(branched[2](tp)) <: Array{Float32} + #ISSUE: what am I really looking for here? + + #Check behaviors of the + @test branched[1](st) == state_to_tuple(st) #Evaluation is of the wrong approach + + + +end #branch generator + +end #overall testset \ No newline at end of file diff --git a/julia_code/Module/workspace.code-workspace b/julia_code/Module/workspace.code-workspace new file mode 100644 index 0000000..bab1b7f --- /dev/null +++ b/julia_code/Module/workspace.code-workspace @@ -0,0 +1,8 @@ +{ + "folders": [ + { + "path": ".." + } + ], + "settings": {} +} \ No newline at end of file