Almost completed redesign of the system. Needs tested.

julia_code/BellmanResidual_Operators.jl

@@ -32,6 +32,9 @@ Three key functions
  - Constellation Satellite Stocks Evolution
 """
 
+# ╔═╡ f8779824-9d59-4458-926f-10beb3cd3866
+abstract type TransitionRule end
+
 # ╔═╡ 9b1c2f62-ec19-4667-8f93-f8b52750e317
 	md"""
 	### Parameterize the physical world
@@ -62,22 +65,15 @@ begin
 		loss_param = 2e-3;
 		loss_weights = loss_param*(ones(N_constellations,N_constellations) - LinearAlgebra.I);
 		
-		#orbital decay rate
-		decay_param = 0.01;
-		
-		#debris generation parameters
-		autocatalysis_param = 0.001;
-		satellite_loss_debris_rate = 5.0;
-		launch_debris_rate = 0.05;
 
 		#Todo, wrap physical model as a struct with the parameters
 		bm = BasicModel(
 		    loss_param
 		    ,loss_weights
-		    ,decay_param
+		    ,0.01
-		    ,autocatalysis_param
+		    ,0.001
-		    ,satellite_loss_debris_rate
+		    ,5.0
-		    ,launch_debris_rate
+		    ,0.05
 		)
 	end
 	
@@ -98,36 +94,51 @@ function survival(
 end
 
 # ╔═╡ 152f3a3c-a565-41bb-8e59-6ab0d2315ffb
+begin
 	#debris evolution 
-function H(
+	struct BasicDebrisEvolution <: TransitionRule
+		physical_model::BasicModel
+	end
+	function (self::BasicDebrisEvolution)(
 			stocks::Array{Float32}
 			,debris::Array{Float32}
 			,launches::Array{Float32}
-		, physical_model::BasicModel
 		)
 	    #get changes in debris from natural dynamics
-    natural_debris_dynamics = (1-physical_model.decay_rate+physical_model.autocatalysis_rate) * debris 
+	    natural_debris_dynamics = (1-self.physical_model.decay_rate + self.physical_model.autocatalysis_rate) * debris 
 	    
 	    #get changes in debris from satellite loss
-    satellite_loss_debris = physical_model.satellite_collision_debris_ratio * (1 .- survival(stocks,debris,physical_model))'*stocks 
+	    satellite_loss_debris = self.physical_model.satellite_collision_debris_ratio * (1 .- survival(stocks,debris,self.physical_model))'*stocks 
 	    
 	    #get changes in debris from launches
-    launch_debris = physical_model.launch_debris_ratio*sum(launches)
+	    launch_debris = self.physical_model.launch_debris_ratio*sum(launches)
 	    
 	    #return total debris level
 	    return natural_debris_dynamics .+ satellite_loss_debris .+ launch_debris
 	end
+end
 
 # ╔═╡ 25ac9438-2b1d-4f6b-9ff1-1695e1d52b51
+begin
 	#stock update rules
-function G(
+	struct BasicStockEvolution <: TransitionRule
+		physical_model::BasicModel
+	end
+	function (self::BasicStockEvolution)(
 			stocks::Array{Float32}
 			,debris::Array{Float32}
 			,launches::Array{Float32}
-		, physical_model::BasicModel
+			
 	)
-    return LinearAlgebra.diagm(survival(stocks,debris,physical_model) .- physical_model.decay_rate)*stocks + launches
+	    return LinearAlgebra.diagm(survival(stocks,debris,self.physical_model) .- self.physical_model.decay_rate)*stocks + launches
 	end
+end
+
+# ╔═╡ 6bf01e30-d054-4a98-9e0d-58d7da4e5524
+begin
+	H = BasicDebrisEvolution(bm)
+	G = BasicStockEvolution(bm)
+end;
 
 # ╔═╡ 29ff1777-d276-4e8f-8582-4ca191f2e2ff
 md"""
@@ -178,35 +189,8 @@ function planner_policy_function_generator(number_params=32)
 
 end
 
-# ╔═╡ 95bfc9d8-8427-41d6-9f0f-f155296eef91
+# ╔═╡ ae6681e0-c796-4145-a313-75f74b4993ad
-#I don't think that these are going to be needed
+#add branch generator
-begin
-	#= CUSTOM LAYERS
-
-	=#
-	#Custom passthrough layer
-	passthrough(x::Array{Float32}) = x
-
-
-	# custom split layer
-	struct Split{T}
-	  paths::T
-	end
-	Split(paths...) = Split(paths)
-	Flux.@functor Split
-	(self::Split)(x::AbstractArray{Float32}) = tuple(map(f -> f(x), self.paths))
-
-	
-	### TESTING ###
-
-
-	#multiple branches
-	Flux.Parallel(vcat,
-		passthrough, passthrough, passthrough
-	)(([1f0],[2f0,3],[4f0]))
-
-
-end
 
 # ╔═╡ bbca5143-f314-40ea-a20e-8a043272e362
 md"""
@@ -214,58 +198,42 @@ md"""
 """
 
 # ╔═╡ 340da189-f443-4376-a82d-7699a21ab7a2
-abstract type EconomicParameters end
+abstract type EconomicModel end
 
 # ╔═╡ 206ac4cc-5102-4381-ad8a-777b02dc4d5a
 begin #basic linear model
-	struct EconModel1 <: EconomicParameters
+	struct EconModel1 <: EconomicModel
 		β::Float32
 		payoff_array::Array{Float32}
 		policy_costs::Array{Float32}
 	end
-
+	function (em::EconModel1)(
-	function payoff1(
 		s::Vector{Float32}
 		,d::Vector{Float32}
 		,a::Vector{Float32}
-	,em::EconModel1
 		)
 		return em.payoff_array*s - em.policy_costs*a
 	end
-
 end
 
 # ╔═╡ eebb8706-a431-4fd1-b7a5-40f07a63d5cb
 begin #basic CES model
-		struct CESParams <: EconomicParameters
+	struct CESParams <: EconomicModel
 		β::Float32
 		r::Float32 #elasticity of subsititution
 		payoff_array::Array{Float32}
 		policy_costs::Array{Float32}
 		debris_costs::Array{Float32}
 	end
-	function CES_with_debris(
+	function (em::CESParams)(
 		s::Vector{Float32}
 		,d::Vector{Float32}
 		,a::Vector{Float32}
-		,em::CESParams
 		)
 		return (em.payoff_array*(s.^em.r) - em.debris_costs*(d.^em.r)).^(1/em.r) - em.policy_costs*a
 	end
 end
 
-# ╔═╡ f8d582cb-10cf-4c72-8127-787f662e0567
-#=
-This struct organizes information about a given constellation operator
-=#
-struct ConstellationOperator
-	payoff_fn::Function
-	econ_params::EconomicParameters
-	value::Flux.Chain
-	#policy_params::Flux.Params #cutting this for now
-end
-#TODO: create a function that takes this struct and checks backprop
-
 # ╔═╡ b433a7ec-8264-48d6-8b95-53d2ec4bad05
 md"""
 ### examples of parameter models
@@ -317,12 +285,19 @@ md"""
 # Loss function specification
 """
 
+# ╔═╡ 9d4a668a-23e3-4f36-86f4-60e242caee3b
+begin #testing data
+	s1 = ones(Float32,N_constellations)
+	d1 = ones(Float32,N_debris)
+end
+
 # ╔═╡ 41271ab4-1ec7-431f-9efb-0f7c3da2d8b4
 #Constellation operator loss function
 function Ξ(
 	s::Vector{Float32}
 	,d::Vector{Float32}
-	,physical_model::PhysicalParameters
+	,debris_transition::BasicDebrisEvolution
+	,stocks_transition::BasicStockEvolution
 	,co::ConstellationOperator
 	,policy::Flux.Chain
 )
@@ -330,8 +305,8 @@ function Ξ(
 	a = policy((s,d))
 
 	#get updated stocks and debris
-	s′ = G(s,d,a,physical_model)
+	s′ = stocks_transition(s,d,a)
-	d′ = H(s,d,a,physical_model)
+	d′ = debris_transition(s,d,a)
 
 	
 	bellman_residuals = co.value((s,d)) - co.payoff_fn(s,d,a,co.econ_params) - co.econ_params.β*co.value((s′,d′))
@@ -341,70 +316,96 @@ function Ξ(
 	return sum(bellman_residuals.^2, maximization_condition)
 end
 
-# ╔═╡ 9d4a668a-23e3-4f36-86f4-60e242caee3b
-begin #testing data
-	s1 = ones(Float32,N_constellations)
-	d1 = ones(Float32,N_debris)
-end
-
-# ╔═╡ cd55e232-493d-4849-8bd7-b0ba85e21bab
-md"""
-# Start setting things up
-"""
-
-# ╔═╡ 08ef7fbf-005b-40b5-acde-f42750c04cd3
-begin
-	const tops = [ 
-		ConstellationOperator(payoff1,em2_a,value_function_generator())
-		,ConstellationOperator(payoff1,em2_b,value_function_generator())
-		,ConstellationOperator(payoff1,em2_c,value_function_generator())
-		,ConstellationOperator(payoff1,em2_d,value_function_generator())
-	]
-
-	#sanity Check time
-	@assert length(tops) == N_constellations "Mismatch in predetermined number of constellations and the number of operators initialized"
-end
-
-# ╔═╡ 43b99708-0052-4b78-886c-92ac2b532f29
-begin #testing
-	Ξ(s1,d1,bm,tops[1],collected_policies)
-end
-
 # ╔═╡ 89258b2e-cee5-4cbd-a42e-21301b7a2549
 begin
 	#=
+	This struct organizes information about a given constellation operator
 	Used to provide an interable loss function for training
+
 	=#
-	struct OperatorLoss
+	struct ConstellationOperatorLoss
-		physical_model::PhysicalParameters
+		#econ model describing operator
-		co::ConstellationOperator
+		econ_model::EconomicModel
-		collected_policies::Flux.Chain
+		#Operator's value and policy functions, as well as which parameters the operator can train.
-		policy_params::Flux.Params
+		operator_value_fn::Flux.Chain
+		collected_policies::Flux.Chain #this is held by all operators
+		operator_policy_params::Flux.Params
+		#Transition functions
+		debris_transition::BasicDebrisEvolution
+		stocks_transition::BasicStockEvolution
 	end
-	function (self::OperatorLoss)(
+	# overriding function to calculate operator loss
+	function (operator::ConstellationOperatorLoss)(
 		s::Vector{Float32}
 		,d::Vector{Float32}
+		; #set default policy to the one held by the ConstellationOperator
+		policy::Flux.Chain = operator.collected_policies
 	)
-		Ξ(s ,d ,self.physical_model ,self.co ,self.collected_policies )
+		
+
+		#get actions
+		a = policy((s,d))
+	
+		#get updated stocks and debris
+		s′ = stocks_transition(s ,d ,a)
+		d′ = debris_transition(s ,d ,a)
+	
+		
+		bellman_residuals = operator.operator_value_fn((s,d)) - operator.econ_model(s,d,a) - operator.econ_model.β * operator.operator_value_fn((s′,d′))
+		
+		maximization_condition = - operator.econ_model(s,d,a,co.econ_params) - co.econ_params.β*co.value((s′,d′))
+	
+		return sum(bellman_residuals.^2, maximization_condition)
 	end
 
+
+	
 	#do the same thing with the planner's problem
 	struct PlannerLoss
+		#=
+		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...
+		=#
+		operators::Array{ConstellationOperatorLoss}
+		planner_policy::Flux.Chain
+		planner_params::Flux.Params
 	end
-	function planners_loss(
+	function (pl::PlannerLoss)(
 		s::Vector{Float32}
 		,d::Vector{Float32}
 	) 
 		l = 0.0
-		for (i,co) in enumerate(tops)
+		for co in pl.operators
-			l += Ξ(s,d,bm,tops,i)
+			l += co(s ,d ,pl.planner_policy)
 		end
 		return l
 	end
+
+	#TODO: Training Functions as routines
+end
+
+# ╔═╡ cd55e232-493d-4849-8bd7-b0ba85e21bab
+md"""
+# Start setting things up
+"""
+
+# ╔═╡ 08ef7fbf-005b-40b5-acde-f42750c04cd3
+begin
+	const tops = [ 
+		ConstellationOperator(payoff1,em2_a,value_function_generator())
+		,ConstellationOperator(payoff1,em2_b,value_function_generator())
+		,ConstellationOperator(payoff1,em2_c,value_function_generator())
+		,ConstellationOperator(payoff1,em2_d,value_function_generator())
+	]
+
+	#sanity Check time
+	@assert length(tops) == N_constellations "Mismatch in predetermined number of constellations and the number of operators initialized"
 end
 
 # ╔═╡ fb6aacff-c42d-4ec1-88cb-5ce1b2e8874f
-policy = planner_policy_function_generator();
+planner_policy = planner_policy_function_generator();
 
 # ╔═╡ 5abebc1a-370c-4f5f-8826-dc0b143d5166
 md"""
@@ -1045,28 +1046,28 @@ uuid = "3f19e933-33d8-53b3-aaab-bd5110c3b7a0"
 # ╟─9fa41b7c-1923-4c1e-bfc6-20ce4a1a2ede
 # ╟─0d209540-29ad-4b2f-9e91-fac2bbee47ff
 # ╟─5b45b29e-f0f4-41e9-91e7-d444687feb4e
-# ╟─152f3a3c-a565-41bb-8e59-6ab0d2315ffb
+# ╠═f8779824-9d59-4458-926f-10beb3cd3866
-# ╟─25ac9438-2b1d-4f6b-9ff1-1695e1d52b51
+# ╠═152f3a3c-a565-41bb-8e59-6ab0d2315ffb
+# ╠═25ac9438-2b1d-4f6b-9ff1-1695e1d52b51
+# ╠═6bf01e30-d054-4a98-9e0d-58d7da4e5524
 # ╠═9b1c2f62-ec19-4667-8f93-f8b52750e317
 # ╟─90446134-4e45-471c-857d-4e165e51937a
 # ╟─29ff1777-d276-4e8f-8582-4ca191f2e2ff
-# ╠═f7aabe43-9a2c-4fe0-8099-c29cdf66566c
+# ╟─f7aabe43-9a2c-4fe0-8099-c29cdf66566c
-# ╠═d816b252-bdca-44ba-ac5c-cb21163a1e9a
+# ╟─d816b252-bdca-44ba-ac5c-cb21163a1e9a
-# ╠═95bfc9d8-8427-41d6-9f0f-f155296eef91
+# ╠═ae6681e0-c796-4145-a313-75f74b4993ad
 # ╠═bbca5143-f314-40ea-a20e-8a043272e362
 # ╠═340da189-f443-4376-a82d-7699a21ab7a2
 # ╠═206ac4cc-5102-4381-ad8a-777b02dc4d5a
 # ╠═eebb8706-a431-4fd1-b7a5-40f07a63d5cb
-# ╠═f8d582cb-10cf-4c72-8127-787f662e0567
 # ╠═b433a7ec-8264-48d6-8b95-53d2ec4bad05
 # ╠═65e0b1fa-d5e1-4ff6-8736-c9d6b5f40150
 # ╠═19ccfc3a-6dbb-4c64-bf03-e2e219ef0efe
 # ╠═dc614254-c211-4552-b985-03020bfc5ab3
 # ╠═4452ab2c-813b-4b7b-8cb3-388ec507a24c
 # ╠═5946daa3-4608-43f3-8933-dd3eb3f4541c
-# ╠═41271ab4-1ec7-431f-9efb-0f7c3da2d8b4
 # ╠═9d4a668a-23e3-4f36-86f4-60e242caee3b
-# ╠═43b99708-0052-4b78-886c-92ac2b532f29
+# ╠═41271ab4-1ec7-431f-9efb-0f7c3da2d8b4
 # ╠═89258b2e-cee5-4cbd-a42e-21301b7a2549
 # ╠═cd55e232-493d-4849-8bd7-b0ba85e21bab
 # ╠═08ef7fbf-005b-40b5-acde-f42750c04cd3