Added policy search capabilities

Project.toml

@@ -24,6 +24,7 @@ SciMLBase = "2"
 julia = "1.10"
 
 [extras]
+DifferentialEquations = "0c46a032-eb83-5123-abaf-570d42b7fbaa"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 NLopt = "76087f3c-5699-56af-9a33-bf431cd00edd"
 NeuralPDE = "315f7962-48a3-4962-8226-d0f33b1235f0"
@@ -35,4 +36,4 @@ SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["SafeTestsets", "Test", "Lux", "Optimization", "OptimizationOptimJL", "OptimizationOptimisers", "NLopt", "Random", "NeuralPDE"]
+test = ["SafeTestsets", "Test", "Lux", "Optimization", "OptimizationOptimJL", "OptimizationOptimisers", "NLopt", "Random", "NeuralPDE", "DifferentialEquations"]

src/NeuralLyapunov.jl

@@ -13,13 +13,15 @@ include("decrease_conditions.jl")
 include("decrease_conditions_RoA_aware.jl")
 include("NeuralLyapunovPDESystem.jl")
 include("local_Lyapunov.jl")
+include("policy_search.jl")
 
 export NeuralLyapunovPDESystem, NumericalNeuralLyapunovFunctions
 export local_Lyapunov
 export NeuralLyapunovSpecification, NeuralLyapunovStructure, UnstructuredNeuralLyapunov,
  NonnegativeNeuralLyapunov, PositiveSemiDefiniteStructure,
  LyapunovMinimizationCondition, StrictlyPositiveDefinite, PositiveSemiDefinite,
  DontCheckNonnegativity, LyapunovDecreaseCondition, AsymptoticDecrease,
- ExponentialDecrease, DontCheckDecrease, RoAAwareDecreaseCondition, make_RoA_aware
+ ExponentialDecrease, DontCheckDecrease, RoAAwareDecreaseCondition, make_RoA_aware,
+ add_policy_search, get_policy
 
 end

src/NeuralLyapunovPDESystem.jl

@@ -35,7 +35,8 @@ function NeuralLyapunovPDESystem(
  fixed_point = zeros(length(lb)),
  p = SciMLBase.NullParameters(),
  state_syms = [],
- parameter_syms = []
+ parameter_syms = [],
+ policy_search::Bool = false
 )::Tuple{PDESystem, Function}
  ########################## Define state symbols ###########################
  state_dim = length(lb)
@@ -78,7 +79,8 @@ function NeuralLyapunovPDESystem(
  fixed_point,
  state,
  params,
- defaults
+ defaults,
+ policy_search
  )
 end
 
@@ -119,7 +121,8 @@ function NeuralLyapunovPDESystem(
  fixed_point = fixed_point,
  p = p,
  state_syms = SciMLBase.variable_symbols(dynamics.sys),
- parameter_syms = p_syms
+ parameter_syms = p_syms,
+ policy_search = false
  )
 end
 
@@ -179,7 +182,8 @@ function NeuralLyapunovPDESystem(
  fixed_point,
  state,
  Num.(parameters(dynamics)),
- ModelingToolkit.get_defaults(dynamics)
+ ModelingToolkit.get_defaults(dynamics),
+ false
  )
 end
 
@@ -191,12 +195,14 @@ function _NeuralLyapunovPDESystem(
  fixed_point,
  state,
  params,
- defaults
+ defaults,
+ policy_search::Bool
 )::Tuple{PDESystem, Function}
  ########################## Unpack specifications ##########################
  structure = spec.structure
  minimzation_condition = spec.minimzation_condition
  decrease_condition = spec.decrease_condition
+ f_call = structure.f_call
 
  ############################# Define domains ##############################
  state_dim = length(lb)
@@ -245,6 +251,10 @@ function _NeuralLyapunovPDESystem(
  push!(bcs, V_sym(fixed_point) ~ 0.0)
  end
 
+ if policy_search
+ append!(bcs, f_call(dynamics, φ, fixed_point, params, 0.0) .~ zeros(state_dim))
+ end
+
  if isempty(eqs) && isempty(bcs)
  error("No training conditions specified.")
  end

src/conditions_specification.jl

@@ -5,18 +5,23 @@ Specifies the structure of the neural Lyapunov function and its derivative.
 
 Allows the user to define the Lyapunov in terms of the neural network to structurally
 enforce Lyapunov conditions.
-`network_dim` is the dimension of the output of the neural network.
 `V(phi::Function, state, fixed_point)` takes in the neural network, the state, and the fixed
 point, and outputs the value of the Lyapunov function at `state`.
 `V̇(phi::Function, J_phi::Function, f::Function, state, params, t, fixed_point)` takes in the
 neural network, jacobian of the neural network, dynamics, state, parameters and time (for
 calling the dynamics, when relevant), and fixed point, and outputs the time derivative of
 the Lyapunov function at `state`.
+`f_call(dynamics::Function, phi::Function, state, p, t)` takes in the dynamics, the neural
+network, the state, the parameters of the dynamics, and time, and outputs the derivative of
+the state; this is useful for making closed-loop dynamics which depend on the neural
+network, such as in the policy search case
+`network_dim` is the dimension of the output of the neural network.
 """
 struct NeuralLyapunovStructure
  V::Function
  ∇V::Function
  V̇::Function
+ f_call::Function
  network_dim::Integer
 end
 

src/policy_search.jl

@@ -0,0 +1,65 @@
+"""
+ add_policy_search(lyapunov_structure, new_dims, control_structure)
+
+Adds dependence on the neural network to the dynamics in a `NeuralLyapunovStructure`
+
+Adds `new_dims` outputs to the neural network and feeds them through `control_structure` to
+calculatethe contribution of the neural network to the dynamics.
+The existing `lyapunov_structure.network_dim` dimensions are used as in `lyapunov_structure`
+to calculate the Lyapunov function.
+
+`lyapunov_structure` should assume in its `V̇` that the dynamics take a form `f(x, p, t)`.
+The returned `NeuralLyapunovStructure` will assume instead `f(x, u, p, t)`, where `u` is the
+contribution from the neural network. Therefore, this structure cannot be used with a
+`NeuralLyapunovPDESystem` method that requires an `ODEFunction`, `ODESystem`, or
+`ODEProblem`.
+"""
+function add_policy_search(
+ lyapunov_structure::NeuralLyapunovStructure,
+ new_dims::Integer;
+ control_structure::Function = identity
+)::NeuralLyapunovStructure
+ let V = lyapunov_structure.V, ∇V = lyapunov_structure.∇V, V̇ = lyapunov_structure.V̇,
+ V_dim = lyapunov_structure.network_dim, nd = new_dims, u = control_structure
+
+ NeuralLyapunovStructure(
+ function (net, state, fixed_point)
+ if length(size(state)) == 1
+ if V_dim == 1
+ V(st -> net(st)[1], state, fixed_point)
+ else
+ V(st -> net(st)[1:V_dim], state, fixed_point)
+ end
+ else
+ V(st -> net(st)[1:V_dim, :], state, fixed_point)
+ end
+ end,
+ function (net, J_net, state, fixed_point)
+ ∇V(st -> net(st)[1:V_dim], st -> J_net(st)[1:V_dim, :], state, fixed_point)
+ end,
+ function (net, J_net, f, state, params, t, fixed_point)
+ V̇(st -> net(st)[1:V_dim], st -> J_net(st)[1:V_dim, :],
+ (st, p, t) -> f(st, u(net(st)[(V_dim + 1):end]), p, t), state, params,
+ t, fixed_point)
+ end,
+ (f, net, state, p, t) -> f(state, u(net(state)[(V_dim + 1):end]), p, t),
+ V_dim + nd
+ )
+ end
+end
+
+function get_policy(
+ phi,
+ θ,
+ network_func::Function,
+ dim_u::Integer;
+ u_func::Function = identity
+)
+ function policy(state::AbstractVector)
+ u_func(network_func(phi, θ, state)[(end - dim_u + 1):end])
+ end
+
+ policy(state::AbstractMatrix) = mapslices(policy, state, dims = [1])
+
+ return policy
+end

src/structure_specification.jl

@@ -3,13 +3,17 @@
 
 Creates a `NeuralLyapunovStructure` where the Lyapunov function is the neural network
 evaluated at the state. This does not structurally enforce any Lyapunov conditions.
+
+Dynamics are assumed to be in `f(state, p, t)` form, as in an `ODEFunction`. For
+`f(state, input, p, t)`, consider using `add_policy_search`.
 """
 function UnstructuredNeuralLyapunov()::NeuralLyapunovStructure
  NeuralLyapunovStructure(
  (net, state, fixed_point) -> net(state),
  (net, grad_net, state, fixed_point) -> grad_net(state),
  (net, grad_net, f, state, params, t, fixed_point) -> grad_net(state) ⋅
  f(state, params, t),
+ (f, net, state, p, t) -> f(state, p, t),
  1
  )
 end
@@ -33,6 +37,9 @@ Lyapunov loss function.
 defaults to `ForwardDiff.gradient`.
 
 The neural network output has dimension `network_dim`.
+
+Dynamics are assumed to be in `f(state, p, t)` form, as in an `ODEFunction`. For
+`f(state, input, p, t)`, consider using `add_policy_search`.
 """
 function NonnegativeNeuralLyapunov(
  network_dim::Integer;
@@ -50,6 +57,7 @@ function NonnegativeNeuralLyapunov(
  (net, J_net, f, state, params, t, fixed_point) -> 2 *
  dot(
  net(state), J_net(state), f(state, params, t)),
+ (f, net, state, p, t) -> f(state, p, t),
  network_dim
  )
  else
@@ -68,6 +76,7 @@ function NonnegativeNeuralLyapunov(
  J_net(state),
  f(state, params, t)
  ) + δ * grad_pos_def(state, fixed_point) ⋅ f(state, params, t),
+ (f, net, state, p, t) -> f(state, p, t),
  network_dim
  )
  end
@@ -94,6 +103,9 @@ gradient of `non_neg(net, state, fixed_point)` with respect to `state` at `state
 defaults to `ForwardDiff.gradient`.
 
 The neural network output has dimension `network_dim`.
+
+Dynamics are assumed to be in `f(state, p, t)` form, as in an `ODEFunction`. For
+`f(state, input, p, t)`, consider using `add_policy_search`.
 """
 function PositiveSemiDefiniteStructure(
  network_dim::Integer;
@@ -133,6 +145,7 @@ function PositiveSemiDefiniteStructure(
  (f(state, params, t) ⋅
  grad_non_neg(
  net, J_net, state, fixed_point)),
+ (f, net, state, p, t) -> f(state, p, t),
  network_dim
  )
 end