diff --git a/Project.toml b/Project.toml
index fe72351..e950ba6 100644
--- a/Project.toml
+++ b/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"]
diff --git a/src/NeuralLyapunov.jl b/src/NeuralLyapunov.jl
index 78171b2..b6e41b7 100644
--- a/src/NeuralLyapunov.jl
+++ b/src/NeuralLyapunov.jl
@@ -13,6 +13,7 @@ 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
@@ -20,6 +21,7 @@ export NeuralLyapunovSpecification, NeuralLyapunovStructure, UnstructuredNeuralL
        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
diff --git a/src/NeuralLyapunovPDESystem.jl b/src/NeuralLyapunovPDESystem.jl
index 349ae2d..81a059a 100644
--- a/src/NeuralLyapunovPDESystem.jl
+++ b/src/NeuralLyapunovPDESystem.jl
@@ -1,55 +1,58 @@
 """
-    NeuralLyapunovPDESystem(dynamics, lb, ub, spec; fixed_point, ps)
+    NeuralLyapunovPDESystem(dynamics::ODESystem, bounds, spec; <keyword_arguments>)
+    NeuralLyapunovPDESystem(dynamics::Function, lb, ub, spec; <keyword_arguments>)
 
-Constructs a ModelingToolkit `PDESystem` to train a neural Lyapunov function
-
-Returns the `PDESystem` and a function representing the neural network, which operates
-columnwise.
-
-The neural Lyapunov function will only be trained for `{ x : lb .≤ x .≤ ub }`. The Lyapunov
-function will be for the dynamical system represented by `dynamics`. If `dynamics` is an
-`ODEProblem` or `ODEFunction`, then the corresponding ODE; if `dynamics` is a function, then
-the ODE is `ẋ = dynamics(x, p, t)`. This ODE should not depend on `t` (time `t=0.0` alone
-will be used) and should have a fixed point at `x = fixed_point`. The particular Lyapunov
-conditions to be used and structure of the neural Lyapunov function are specified through
-`spec`, which is a `NeuralLyapunovSpecification`.
+Construct and return a `PDESystem` representing the specified neural Lyapunov problem, along
+with a function representing the neural network.
 
 The returned neural network function takes three inputs: the neural network structure `phi`,
-the trained network parameters, and a matrix of inputs to operate on columnwise.
-
-If `dynamics` requires parameters, their values can be supplied through the Vector `p`, or
-through the parameters of `dynamics` if `dynamics isa ODEProblem` (in which case, let
-the other be `SciMLBase.NullParameters()`). If `dynamics` is an `ODEFunction` and
-`dynamics.paramsyms` is defined, then `p` should have the same order.
+the trained network parameters, and a matrix of inputs, then operates columnwise on the
+inputs.
+
+# Arguments
+- `dynamics`: the dynamical system being analyzed, represented as an `ODESystem` or the
+        function `f` such that `ẋ = f(x[, u], p, t)`; either way, the ODE should not depend
+        on time and only `t = 0.0` will be used
+- `bounds`: an array of domains, defining the training domain by bounding the states (and
+        derivatives, when applicable) of `dynamics`; only used when `dynamics isa
+        ODESystem`, otherwise use `lb` and `ub`.
+- `lb` and `ub`: the training domain will be ``[lb_1, ub_1]×[lb_2, ub_2]×...``; not used
+        when `dynamics isa ODESystem`, then use `bounds`.
+- `spec::NeuralLyapunovSpecification`: defines the Lyapunov function structure, as well as
+        the minimization and decrease conditions.
+- `fixed_point`: the equilibrium being analyzed; defaults to the origin.
+- `p`: the values of the parameters of the dynamical system being analyzed; defaults to
+        `SciMLBase.NullParameters()`; not used when `dynamics isa ODESystem`, then use the
+        default parameter values of `dynamics`.
+- `state_syms`: an array of the `Symbol` representing each state; not used when `dynamics
+        isa ODESystem`, then the symbols from `dynamics` are used; if `dynamics isa
+        ODEFunction`, symbols stored there are used, unless overridden here; if not provided
+        here and cannot be inferred, `[:state1, :state2, ...]` will be used.
+- `parameter_syms`: an array of the `Symbol` representing each parameter; not used when
+        `dynamics isa ODESystem`, then the symbols from `dynamics` are used; if `dynamics
+        isa ODEFunction`, symbols stored there are used, unless overridden here; if not
+        provided here and cannot be inferred, `[:param1, :param2, ...]` will be used.
+- `policy_search::Bool`: whether or not to include a loss term enforcing `fixed_point` to
+        actually be a fixed point; defaults to `false`; only used when `dynamics isa
+        Function && !(dynamics isa ODEFunction)`; when `dynamics isa ODEFunction`,
+        `policy_search` must be `false`, so should not be supplied; when `dynamics isa
+        ODESystem`, value inferred by the presence of unbound inputs.
 """
 function NeuralLyapunovPDESystem(
-        dynamics::ODEFunction,
+        dynamics::Function,
         lb,
         ub,
         spec::NeuralLyapunovSpecification;
         fixed_point = zeros(length(lb)),
-        p = SciMLBase.NullParameters()
+        p = SciMLBase.NullParameters(),
+        state_syms = [],
+        parameter_syms = [],
+        policy_search::Bool = false
 )::Tuple{PDESystem, Function}
-    if dynamics.mass_matrix !== I
-        throw(ErrorException("DAEs are not supported at this time."))
-    end
-
-    if dynamics.sys isa ODESystem
-        return NeuralLyapunovPDESystem(
-            dynamics.sys,
-            lb,
-            ub,
-            spec;
-            fixed_point = fixed_point,
-            p = p
-        )
-    end
-
     ########################## Define state symbols ###########################
     state_dim = length(lb)
 
     # Define state symbols, if not already defined
-    state_syms = SciMLBase.variable_symbols(dynamics.sys)
     state_syms = if isempty(state_syms)
         [Symbol(:state, i) for i in 1:state_dim]
     else
@@ -62,10 +65,10 @@ function NeuralLyapunovPDESystem(
     param_syms = if p == SciMLBase.NullParameters()
         []
     else
-        if isnothing(dynamics.sys.parameters)
+        if isempty(parameter_syms)
             [Symbol(:param, i) for i in 1:length(p)]
         else
-            dynamics.sys.parameters
+            parameter_syms
         end
     end
 
@@ -78,134 +81,146 @@ function NeuralLyapunovPDESystem(
         Dict([param => param_val for (param, param_val) in zip(params, p)])
     end
 
+    ############################# Define domains ##############################
+    domains = [state[i] ∈ (lb[i], ub[i]) for i in 1:state_dim]
+
     ########################### Construct PDESystem ###########################
     return _NeuralLyapunovPDESystem(
         dynamics,
-        lb,
-        ub,
+        domains,
         spec,
         fixed_point,
         state,
         params,
-        defaults
+        defaults,
+        policy_search
     )
 end
 
 function NeuralLyapunovPDESystem(
-        dynamics::Function,
+        dynamics::ODEFunction,
         lb,
         ub,
         spec::NeuralLyapunovSpecification;
         fixed_point = zeros(length(lb)),
         p = SciMLBase.NullParameters()
 )::Tuple{PDESystem, Function}
-    return NeuralLyapunovPDESystem(
-        ODEFunction(dynamics),
-        lb,
-        ub,
-        spec;
-        fixed_point = fixed_point,
-        p = p
-    )
-end
+    if dynamics.mass_matrix !== I
+        throw(ErrorException("DAEs are not supported at this time. Please supply dynamics" *
+                             " without a mass matrix."))
+    end
 
-function NeuralLyapunovPDESystem(
-        dynamics::ODEProblem,
-        lb,
-        ub,
-        spec::NeuralLyapunovSpecification;
-        fixed_point = zeros(length(lb)),
-        p = SciMLBase.NullParameters()
-)::Tuple{PDESystem, Function}
-    f = dynamics.f
-
-    p = if dynamics.p == SciMLBase.NullParameters()
-        p
-    elseif p == SciMLBase.NullParameters()
-        dynamics.p
-    elseif dynamics.p == p
-        p
+    s_syms, p_syms = if dynamics.sys isa ODESystem
+        s_syms = Symbol.(operation.(states(dynamics.sys)))
+        p_syms = Symbol.(parameters(dynamics.sys))
+        (s_syms, p_syms)
+    elseif dynamics.sys isa SciMLBase.SymbolCache
+        s_syms = SciMLBase.variable_symbols(dynamics.sys)
+        p_syms = if isnothing(dynamics.sys.parameters)
+            []
+        else
+            dynamics.sys.parameters
+        end
+        (s_syms, p_syms)
     else
-        throw(ErrorException("Conflicting parameter definitions. Please define parameters" *
-                             " only through p or dynamics.p; the other should be " *
-                             "SciMLBase.NullParameters()"))
+        ([], [])
     end
 
     return NeuralLyapunovPDESystem(
-        f,
+        dynamics.f,
         lb,
         ub,
         spec;
         fixed_point = fixed_point,
-        p = p
+        p = p,
+        state_syms = s_syms,
+        parameter_syms = p_syms,
+        policy_search = false
     )
 end
 
 function NeuralLyapunovPDESystem(
         dynamics::ODESystem,
-        lb,
-        ub,
+        bounds,
         spec::NeuralLyapunovSpecification;
-        fixed_point = zeros(length(lb))
+        fixed_point = zeros(length(bounds))
 )::Tuple{PDESystem, Function}
+    ######################### Check for policy search #########################
+    f, x, p, policy_search = if isempty(ModelingToolkit.unbound_inputs(dynamics))
+        (ODEFunction(dynamics), states(dynamics), parameters(dynamics), false)
+    else
+        (f, _), x, p = ModelingToolkit.generate_control_function(dynamics; simplify = true)
+        (f, x, p, true)
+    end
+
     ########################## Define state symbols ###########################
-    state = states(dynamics)
     # States should all be functions of time, but we just want the symbol
     # e.g., if the state is ω(t), we just want ω
-    state = map(st -> istree(st) ? operation(st) : st, state)
-    state_syms = Symbol.(state)
+    _state = operation.(x)
+    state_syms = Symbol.(_state)
     state = [first(@parameters $s) for s in state_syms]
 
+    ###################### Remove derivatives in domains ######################
+    domains = map(
+        d -> Num(operation(Symbolics.diff2term(Symbolics.value(d.variables)))) ∈ d.domain,
+        bounds
+    )
+    domain_vars = map(d -> d.variables, domains)
+    if Set(_state) != Set(domain_vars)
+        error("Domain variables from `domains` do not match those extracted from " *
+              "`dynamics`. Got $_state from `dynamics` and $domain_vars from `domains`.")
+    end
+
     ########################### Construct PDESystem ###########################
     _NeuralLyapunovPDESystem(
-        ODEFunction(dynamics),
-        lb,
-        ub,
+        f,
+        domains,
         spec,
         fixed_point,
         state,
-        Num.(parameters(dynamics)),
-        ModelingToolkit.get_defaults(dynamics)
+        p,
+        ModelingToolkit.get_defaults(dynamics),
+        policy_search
     )
 end
 
 function _NeuralLyapunovPDESystem(
         dynamics::Function,
-        lb,
-        ub,
+        domains,
         spec::NeuralLyapunovSpecification,
         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
-
-    ############################# Define domains ##############################
-    state_dim = length(lb)
-    domains = [state[i] ∈ (lb[i], ub[i]) for i in 1:state_dim]
+    f_call = structure.f_call
+    state_dim = length(domains)
 
     ################## Define Lyapunov function & derivative ##################
     output_dim = structure.network_dim
-    net_syms = [Symbol(:u, i) for i in 1:output_dim]
+    net_syms = [Symbol(:φ, i) for i in 1:output_dim]
     net = [first(@variables $s(..)) for s in net_syms]
 
-    # u(x) is the symbolic form of neural network output
-    u(x) = Num.([ui(x...) for ui in net])
+    # φ(x) is the symbolic form of neural network output
+    φ(x) = Num.([φi(x...) for φi in net])
 
     # V_sym(x) is the symobolic form of the Lyapunov function
-    V_sym(x) = structure.V(u, x, fixed_point)
+    V_sym(x) = structure.V(φ, x, fixed_point)
 
     # V̇_sym(x) is the symbolic time derivative of the Lyapunov function
     function V̇_sym(x)
         structure.V̇(
-            u,
-            y -> Symbolics.jacobian(u(y), y),
-            y -> dynamics(y, params, 0.0),
+            φ,
+            y -> Symbolics.jacobian(φ(y), y),
+            dynamics,
             x,
+            params,
+            0.0,
             fixed_point
         )
     end
@@ -225,11 +240,12 @@ function _NeuralLyapunovPDESystem(
 
     bcs = []
 
-    if check_fixed_point(minimzation_condition)
+    if check_minimal_fixed_point(minimzation_condition)
         push!(bcs, V_sym(fixed_point) ~ 0.0)
     end
-    if check_stationary_fixed_point(decrease_condition)
-        push!(bcs, V̇_sym(fixed_point) ~ 0.0)
+
+    if policy_search
+        append!(bcs, f_call(dynamics, φ, fixed_point, params, 0.0) .~ zeros(state_dim))
     end
 
     if isempty(eqs) && isempty(bcs)
@@ -237,7 +253,11 @@ function _NeuralLyapunovPDESystem(
     end
 
     # NeuralPDE requires an equation and a boundary condition, even if they are
-    # trivial like 0.0 == 0.0
+    # trivial like 0.0 == 0.0, so we remove those trivial equations if they showed up
+    # naturally alongside other equations and add them in if we have no other equations
+    eqs = filter(eq -> eq != (0.0 ~ 0.0), eqs)
+    bcs = filter(eq -> eq != (0.0 ~ 0.0), bcs)
+
     if isempty(eqs)
         push!(eqs, 0.0 ~ 0.0)
     end
@@ -251,21 +271,21 @@ function _NeuralLyapunovPDESystem(
         bcs,
         domains,
         state,
-        u(state),
+        φ(state),
         params;
         defaults = defaults
     )
 
     ################### Return PDESystem and neural network ###################
-    # u_func is the numerical form of neural network output
-    function u_func(phi, θ, x)
+    # φ_func is the numerical form of neural network output
+    function φ_func(phi, θ, x)
         reduce(
             vcat,
             Array(phi[i](x, θ.depvar[net_syms[i]])) for i in 1:output_dim
         )
     end
 
-    return lyapunov_pde_system, u_func
+    return lyapunov_pde_system, φ_func
 end
 
 """
@@ -292,7 +312,7 @@ function NumericalNeuralLyapunovFunctions(
         structure::NeuralLyapunovStructure,
         dynamics::Function,
         fixed_point;
-        p,
+        p = SciMLBase.NullParameters(),
         jac = ForwardDiff.jacobian,
         J_net = (_phi, _θ, x) -> jac((y) -> network_func(_phi, _θ, y), x)
 )::Tuple{Function, Function, Function}
@@ -318,8 +338,10 @@ function NumericalNeuralLyapunovFunctions(
         structure.V̇(
             _net_func,
             _J_net,
-            y -> dynamics(y, p, 0.0),
+            dynamics,
             state,
+            p,
+            0.0,
             fixed_point
         )
     end
@@ -350,7 +372,7 @@ function NumericalNeuralLyapunovFunctions(
         V_structure::Function,
         dynamics::Function,
         fixed_point;
-        p = SciMLBase.NullParameters,
+        p = SciMLBase.NullParameters(),
         grad = ForwardDiff.gradient
 )::Tuple{Function, Function, Function}
     # Make network function
diff --git a/src/conditions_specification.jl b/src/conditions_specification.jl
index 8e30305..010e7b8 100644
--- a/src/conditions_specification.jl
+++ b/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, fixed_point)` takes in the neural
-network, the jacobian of the neural network, the dynamics (as a function of the state
-alone), the state, and the fixed point, and outputs the time derivative 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
 
@@ -36,7 +41,7 @@ abstract type AbstractLyapunovMinimizationCondition end
 Represents the decrease condition in a neural Lyapunov problem
 
 All concrete `AbstractLyapunovDecreaseCondition` subtypes should define the
-`check_decrease`, `check_stationary_fixed_point`, and `get_decrease_condition` functions.
+`check_decrease` and `get_decrease_condition` functions.
 """
 abstract type AbstractLyapunovDecreaseCondition end
 
@@ -63,14 +68,14 @@ function check_nonnegativity(cond::AbstractLyapunovMinimizationCondition)::Bool
 end
 
 """
-    check_fixed_point(cond::AbstractLyapunovMinimizationCondition)
+    check_minimal_fixed_point(cond::AbstractLyapunovMinimizationCondition)
 
 `true` if `cond` specifies training for the Lyapunov function to equal zero at the
 fixed point, `false` if `cond` specifies no training to meet this condition.
 """
-function check_fixed_point(cond::AbstractLyapunovMinimizationCondition)::Bool
-    error("check_fixed_point not implemented for AbstractLyapunovMinimizationCondition of" *
-          " type $(typeof(cond))")
+function check_minimal_fixed_point(cond::AbstractLyapunovMinimizationCondition)::Bool
+    error("check_minimal_fixed_point not implemented for " *
+          "AbstractLyapunovMinimizationCondition of type $(typeof(cond))")
 end
 
 """
@@ -95,17 +100,6 @@ function check_decrease(cond::AbstractLyapunovDecreaseCondition)::Bool
           string(typeof(cond)))
 end
 
-"""
-    check_stationary_fixed_point(cond::AbstractLyapunovDecreaseCondition)
-
-`true` if `cond` specifies training for the Lyapunov function not to change at the
-fixed point, `false` if `cond` specifies no training to meet this condition.
-"""
-function check_stationary_fixed_point(cond::AbstractLyapunovDecreaseCondition)::Bool
-    error("check_fixed_point not implemented for AbstractLyapunovDecreaseCondition of " *
-          "type $(typeof(cond))")
-end
-
 """
     get_decrease_condition(cond::AbstractLyapunovDecreaseCondition)
 
diff --git a/src/decrease_conditions.jl b/src/decrease_conditions.jl
index 48353d4..f173092 100644
--- a/src/decrease_conditions.jl
+++ b/src/decrease_conditions.jl
@@ -1,13 +1,11 @@
 """
-    LyapunovDecreaseCondition(check_decrease, decrease, strength, relu, check_fixed_point)
+    LyapunovDecreaseCondition(check_decrease, decrease, strength, relu)
 
 Specifies the form of the Lyapunov conditions to be used; if `check_decrease`, training will
 enforce `decrease(V, dVdt) ≤ strength(state, fixed_point)`.
 
 The inequality will be approximated by the equation
     `relu(decrease(V, dVdt) - strength(state, fixed_point)) = 0.0`.
-If `check_fixed_point` is `false`, then training assumes `dVdt(fixed_point) = 0`, but
-if `check_fixed_point` is `true`, then training will enforce `dVdt(fixed_point) = 0`.
 
 If the dynamics truly have a fixed point at `fixed_point` and `dVdt` has been defined
 properly in terms of the dynamics, then `dVdt(fixed_point)` will be `0` and there is no need
@@ -30,17 +28,12 @@ struct LyapunovDecreaseCondition <: AbstractLyapunovDecreaseCondition
     decrease::Function
     strength::Function
     relu::Function
-    check_fixed_point::Bool
 end
 
 function check_decrease(cond::LyapunovDecreaseCondition)::Bool
     cond.check_decrease
 end
 
-function check_stationary_fixed_point(cond::LyapunovDecreaseCondition)::Bool
-    cond.check_fixed_point
-end
-
 function get_decrease_condition(cond::LyapunovDecreaseCondition)
     if cond.check_decrease
         return (V, dVdt, x, fixed_point) -> cond.relu(
@@ -52,7 +45,7 @@ function get_decrease_condition(cond::LyapunovDecreaseCondition)
 end
 
 """
-    AsymptoticDecrease(strict; check_fixed_point, C)
+    AsymptoticDecrease(; strict, C, relu)
 
 Constructs a `LyapunovDecreaseCondition` corresponding to asymptotic decrease.
 
@@ -63,7 +56,6 @@ The inequality is represented by `a ≥ b` <==> `relu(b-a) = 0.0`.
 """
 function AsymptoticDecrease(;
         strict::Bool = false,
-        check_fixed_point::Bool = false,
         C::Real = 1e-6,
         relu = (t) -> max(0.0, t)
 )::LyapunovDecreaseCondition
@@ -77,13 +69,12 @@ function AsymptoticDecrease(;
         true,
         (V, dVdt) -> dVdt,
         strength,
-        relu,
-        check_fixed_point
+        relu
     )
 end
 
 """
-    ExponentialDecrease(k, strict; check_fixed_point, C)
+    ExponentialDecrease(k; strict, C, relu)
 
 Constructs a `LyapunovDecreaseCondition` corresponding to exponential decrease of rate `k`.
 
@@ -95,7 +86,6 @@ The inequality is represented by `a ≥ b` <==> `relu(b-a) = 0.0`.
 function ExponentialDecrease(
         k::Real;
         strict::Bool = false,
-        check_fixed_point::Bool = false,
         C::Real = 1e-6,
         relu = (t) -> max(0.0, t)
 )::LyapunovDecreaseCondition
@@ -109,26 +99,22 @@ function ExponentialDecrease(
         true,
         (V, dVdt) -> dVdt + k * V,
         strength,
-        relu,
-        check_fixed_point
+        relu
     )
 end
 
 """
-    DontCheckDecrease(check_fixed_point = false)
+    DontCheckDecrease()
 
 Constructs a `LyapunovDecreaseCondition` which represents not checking for
 decrease of the Lyapunov function along system trajectories. This is appropriate
 in cases when the Lyapunov decrease condition has been structurally enforced.
-
-It is still possible to check for `dV/dt = 0` at `fixed_point`, even in this case.
 """
-function DontCheckDecrease(check_fixed_point::Bool = false)::LyapunovDecreaseCondition
+function DontCheckDecrease()::LyapunovDecreaseCondition
     return LyapunovDecreaseCondition(
         false,
         (V, dVdt) -> 0.0,
         (state, fixed_point) -> 0.0,
-        (t) -> 0.0,
-        check_fixed_point
+        (t) -> 0.0
     )
 end
diff --git a/src/decrease_conditions_RoA_aware.jl b/src/decrease_conditions_RoA_aware.jl
index d6a9de3..fba7e1f 100644
--- a/src/decrease_conditions_RoA_aware.jl
+++ b/src/decrease_conditions_RoA_aware.jl
@@ -1,6 +1,6 @@
 """
-    RoAAwareDecreaseCondition(check_decrease, decrease, strength, relu, check_fixed_point,
-                              ρ, out_of_RoA_penalty)
+    RoAAwareDecreaseCondition(check_decrease, decrease, strength, relu, ρ,
+                              out_of_RoA_penalty)
 
 Specifies the form of the Lyapunov conditions to be used, training for a region of
 attraction estimate of `{ x : V(x) ≤ ρ }`
@@ -12,9 +12,6 @@ and will instead apply
 
 The inequality will be approximated by the equation
     `relu(decrease(V, dVdt) - strength(state, fixed_point)) = 0.0`.
-If `check_fixed_point` is `false`, then training assumes `dVdt(fixed_point) = 0`, but
-if `check_fixed_point` is `true`, then training will attempt to enforce
-`dVdt(fixed_point) = 0`.
 
 If the dynamics truly have a fixed point at `fixed_point` and `dVdt` has been defined
 properly in terms of the dynamics, then `dVdt(fixed_point)` will be `0` and there is no need
@@ -44,7 +41,6 @@ struct RoAAwareDecreaseCondition <: AbstractLyapunovDecreaseCondition
     decrease::Function
     strength::Function
     relu::Function
-    check_fixed_point::Bool
     ρ::Real
     out_of_RoA_penalty::Function
 end
@@ -53,10 +49,6 @@ function check_decrease(cond::RoAAwareDecreaseCondition)::Bool
     cond.check_decrease
 end
 
-function check_stationary_fixed_point(cond::RoAAwareDecreaseCondition)::Bool
-    cond.check_fixed_point
-end
-
 function get_decrease_condition(cond::RoAAwareDecreaseCondition)
     if cond.check_decrease
         return function (V, dVdt, x, fixed_point)
@@ -92,7 +84,6 @@ function make_RoA_aware(
         cond.decrease,
         cond.strength,
         cond.relu,
-        cond.check_fixed_point,
         ρ,
         out_of_RoA_penalty
     )
diff --git a/src/minimization_conditions.jl b/src/minimization_conditions.jl
index f19b53b..e11af3e 100644
--- a/src/minimization_conditions.jl
+++ b/src/minimization_conditions.jl
@@ -38,7 +38,7 @@ function check_nonnegativity(cond::LyapunovMinimizationCondition)::Bool
     cond.check_nonnegativity
 end
 
-function check_fixed_point(cond::LyapunovMinimizationCondition)::Bool
+function check_minimal_fixed_point(cond::LyapunovMinimizationCondition)::Bool
     cond.check_fixed_point
 end
 
diff --git a/src/policy_search.jl b/src/policy_search.jl
new file mode 100644
index 0000000..80b055e
--- /dev/null
+++ b/src/policy_search.jl
@@ -0,0 +1,78 @@
+"""
+    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
+
+"""
+    get_policy(phi, θ, network_func, dim; control_structure)
+
+Returns the control policy as a function of the state
+
+The returned function can operate on a state vector or columnwise on a matrix of state
+vectors.
+
+`phi` is the neural network with parameters `θ`. `network_func` is an output of
+`NeuralLyapunovPDESystem`.
+The control policy is `control_structure` composed with the last `dim` outputs of the
+neural network, as set up by `add_policy_search`.
+"""
+function get_policy(
+        phi,
+        θ,
+        network_func::Function,
+        dim::Integer;
+        control_structure::Function = identity
+)
+    function policy(state::AbstractVector)
+        control_structure(network_func(phi, θ, state)[(end - dim + 1):end])
+    end
+
+    policy(state::AbstractMatrix) = mapslices(policy, state, dims = [1])
+
+    return policy
+end
diff --git a/src/structure_specification.jl b/src/structure_specification.jl
index e225fb4..b1f6af5 100644
--- a/src/structure_specification.jl
+++ b/src/structure_specification.jl
@@ -3,12 +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, fixed_point) -> grad_net(state) ⋅ f(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
@@ -32,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;
@@ -46,8 +54,10 @@ function NonnegativeNeuralLyapunov(
         NeuralLyapunovStructure(
             (net, state, fixed_point) -> net(state) ⋅ net(state),
             (net, J_net, state, fixed_point) -> 2 * transpose(net(state)) * J_net(state),
-            (net, J_net, f, state, fixed_point) -> 2 *
-                                                   dot(net(state), J_net(state), f(state)),
+            (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
@@ -61,11 +71,12 @@ function NonnegativeNeuralLyapunov(
                                          δ * pos_def(state, fixed_point),
             (net, J_net, state, fixed_point) -> 2 * transpose(net(state)) * J_net(state) +
                                                 δ * grad_pos_def(state, fixed_point),
-            (net, J_net, f, state, fixed_point) -> 2 * dot(
+            (net, J_net, f, state, params, t, fixed_point) -> 2 * dot(
                 net(state),
                 J_net(state),
-                f(state)
-            ) + δ * grad_pos_def(state, fixed_point) ⋅ f(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
@@ -92,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;
@@ -123,12 +137,15 @@ function PositiveSemiDefiniteStructure(
                                             non_neg(net, state, fixed_point) +
                                             pos_def(state, fixed_point) *
                                             grad_non_neg(net, J_net, state, fixed_point),
-        (net, J_net, f, state, fixed_point) -> (f(state) ⋅
-                                                grad_pos_def(state, fixed_point)) *
-                                               non_neg(net, state, fixed_point) +
-                                               pos_def(state, fixed_point) *
-                                               (f(state) ⋅ grad_non_neg(
+        (net, J_net, f, state, params, t, fixed_point) -> (f(state, params, t) ⋅
+                                                           grad_pos_def(
+            state, fixed_point)) *
+                                                          non_neg(net, state, fixed_point) +
+                                                          pos_def(state, fixed_point) *
+                                                          (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
diff --git a/test/damped_pendulum.jl b/test/damped_pendulum.jl
index 6632021..fcda6e9 100644
--- a/test/damped_pendulum.jl
+++ b/test/damped_pendulum.jl
@@ -29,16 +29,19 @@ eqs = [DDt(θ) + 2ζ * Dt(θ) + ω_0^2 * sin(θ) ~ 0.0]
 )
 
 dynamics = structural_simplify(dynamics)
-
-lb = [-pi, -10.0];
-ub = [pi, 10.0];
+bounds = [
+    θ ∈ (-π, π),
+    Dt(θ) ∈ (-10.0, 10.0)
+]
+lb = [-π, -10.0];
+ub = [π, 10.0];
 p = [defaults[param] for param in parameters(dynamics)]
 
 ####################### Specify neural Lyapunov problem #######################
 
 # Define neural network discretization
 # We use an input layer that is periodic with period 2π with respect to θ
-dim_state = length(lb)
+dim_state = length(bounds)
 dim_hidden = 15
 dim_output = 2
 chain = [Lux.Chain(
@@ -80,10 +83,11 @@ spec = NeuralLyapunovSpecification(
 ############################# Construct PDESystem #############################
 
 pde_system, network_func = NeuralLyapunovPDESystem(
-    dynamics,
+    ODEFunction(dynamics),
     lb,
     ub,
-    spec
+    spec;
+    p = p
 )
 
 ######################## Construct OptimizationProblem ########################
@@ -105,7 +109,7 @@ V_func, V̇_func, ∇V_func = NumericalNeuralLyapunovFunctions(
     network_func,
     structure.V,
     ODEFunction(dynamics),
-    zeros(length(lb));
+    zeros(length(bounds));
     p = p
 )
 
diff --git a/test/inverted_pendulum.jl b/test/inverted_pendulum.jl
new file mode 100644
index 0000000..54dfbf0
--- /dev/null
+++ b/test/inverted_pendulum.jl
@@ -0,0 +1,246 @@
+using LinearAlgebra
+using NeuralPDE, Lux, ModelingToolkit
+using Optimization, OptimizationOptimisers, OptimizationOptimJL, NLopt
+using NeuralLyapunov
+using Random
+using Test
+
+Random.seed!(200)
+
+println("Inverted Pendulum - Policy Search")
+
+######################### Define dynamics and domain ##########################
+
+function open_loop_pendulum_dynamics(x, u, p, t)
+    θ, ω = x
+    ζ, ω_0 = p
+    τ = u[]
+    return [ω
+            -2ζ * ω - ω_0^2 * sin(θ) + τ]
+end
+
+lb = [0.0, -10.0];
+ub = [2π, 10.0];
+upright_equilibrium = [π, 0.0]
+p = [0.5, 1.0]
+state_syms = [:θ, :ω]
+parameter_syms = [:ζ, :ω_0]
+
+####################### Specify neural Lyapunov problem #######################
+
+# Define neural network discretization
+# We use an input layer that is periodic with period 2π with respect to θ
+dim_state = length(lb)
+dim_hidden = 15
+dim_phi = 2
+dim_u = 1
+dim_output = dim_phi + dim_u
+chain = [Lux.Chain(
+             Lux.WrappedFunction(x -> vcat(
+                 transpose(sin.(x[1, :])),
+                 transpose(cos.(x[1, :])),
+                 transpose(x[2, :])
+             )),
+             Dense(3, dim_hidden, tanh),
+             Dense(dim_hidden, dim_hidden, tanh),
+             Dense(dim_hidden, 1, use_bias = false)
+         ) for _ in 1:dim_output]
+
+# Define neural network discretization
+strategy = GridTraining(0.1)
+discretization = PhysicsInformedNN(chain, strategy)
+
+# Define neural Lyapunov structure
+structure = PositiveSemiDefiniteStructure(
+    dim_phi;
+    pos_def = function (state, fixed_point)
+        θ, ω = state
+        θ_eq, ω_eq = fixed_point
+        log(1.0 + (sin(θ) - sin(θ_eq))^2 + (cos(θ) - cos(θ_eq))^2 + (ω - ω_eq)^2)
+    end
+)
+structure = add_policy_search(
+    structure,
+    dim_u
+)
+minimization_condition = DontCheckNonnegativity(check_fixed_point = false)
+
+# Define Lyapunov decrease condition
+decrease_condition = AsymptoticDecrease(strict = true)
+
+# Construct neural Lyapunov specification
+spec = NeuralLyapunovSpecification(
+    structure,
+    minimization_condition,
+    decrease_condition
+)
+
+############################# Construct PDESystem #############################
+
+pde_system, network_func = NeuralLyapunovPDESystem(
+    open_loop_pendulum_dynamics,
+    lb,
+    ub,
+    spec;
+    fixed_point = upright_equilibrium,
+    p = p,
+    state_syms = state_syms,
+    parameter_syms = parameter_syms,
+    policy_search = true
+)
+
+######################## Construct OptimizationProblem ########################
+
+sym_prob = symbolic_discretize(pde_system, discretization)
+prob = discretize(pde_system, discretization)
+
+########################## Solve OptimizationProblem ##########################
+
+res = Optimization.solve(prob, OptimizationOptimisers.Adam(); maxiters = 300)
+prob = Optimization.remake(prob, u0 = res.u)
+res = Optimization.solve(prob, BFGS(); maxiters = 300)
+
+###################### Get numerical numerical functions ######################
+
+V_func, V̇_func, ∇V_func = NumericalNeuralLyapunovFunctions(
+    discretization.phi,
+    res.u,
+    network_func,
+    structure,
+    open_loop_pendulum_dynamics,
+    upright_equilibrium;
+    p = p
+)
+
+u = get_policy(discretization.phi, res.u, network_func, dim_u)
+
+################################## Simulate ###################################
+
+xs = (-2π):0.02:(2π)
+ys = lb[2]:0.02:ub[2]
+states = Iterators.map(collect, Iterators.product(xs, ys))
+V_predict = vec(V_func(hcat(states...)))
+dVdt_predict = vec(V̇_func(hcat(states...)))
+
+#################################### Tests ####################################
+
+# Network structure should enforce positive definiteness
+@test V_func(upright_equilibrium) == 0.0
+@test min(V_func(upright_equilibrium), minimum(V_predict)) ≥ 0.0
+
+# Network structure should enforce periodicity in θ
+x0 = (ub .- lb) .* rand(2, 100) .+ lb
+@test all(isapprox.(V_func(x0), V_func(x0 .+ [2π, 0.0]); rtol = 1e-3))
+
+# Training should result in a fixed point at the upright equilibrium
+@test all(isapprox.(
+    open_loop_pendulum_dynamics(upright_equilibrium, u(upright_equilibrium), p, 0.0),
+    0.0; atol = 1e-8))
+@test V̇_func(upright_equilibrium) == 0.0
+
+# V̇ should be negative almost everywhere
+@test sum(dVdt_predict .> 0) / length(dVdt_predict) < 1e-3
+
+################################## Simulate ###################################
+
+using DifferentialEquations
+
+closed_loop_dynamics = ODEFunction(
+    (x, p, t) -> open_loop_pendulum_dynamics(x, u(x), p, t);
+    sys = SciMLBase.SymbolCache(state_syms, parameter_syms)
+)
+
+# Starting still at bottom
+downward_equilibrium = zeros(2)
+ode_prob = ODEProblem(closed_loop_dynamics, downward_equilibrium, [0.0, 35.0], p)
+sol = solve(ode_prob, Tsit5())
+# plot(sol)
+
+# Should make it to the top
+θ_end, ω_end = sol.u[end]
+x_end, y_end = sin(θ_end), -cos(θ_end)
+@test all(isapprox.([x_end, y_end, ω_end], [0.0, 1.0, 0.0]; atol = 1e-3))
+
+# Starting at a random point
+x0 = lb .+ rand(2) .* (ub .- lb)
+ode_prob = ODEProblem(closed_loop_dynamics, x0, [0.0, 20.0], p)
+sol = solve(ode_prob, Tsit5())
+# plot(sol)
+
+# Should make it to the top
+θ_end, ω_end = sol.u[end]
+x_end, y_end = sin(θ_end), -cos(θ_end)
+@test all(isapprox.([x_end, y_end, ω_end], [0.0, 1.0, 0.0]; atol = 1e-3))
+
+#=
+# Print statistics
+println("V(π, 0) = ", V_func(upright_equilibrium))
+println(
+    "f([π, 0], u([π, 0])) = ",
+    open_loop_pendulum_dynamics(upright_equilibrium, u(upright_equilibrium), p, 0.0)
+)
+println(
+    "V ∋ [",
+    min(V_func(upright_equilibrium),
+    minimum(V_predict)),
+    ", ",
+    maximum(V_predict),
+    "]"
+)
+println(
+    "V̇ ∋ [",
+    minimum(dVdt_predict),
+    ", ",
+    max(V̇_func(upright_equilibrium), maximum(dVdt_predict)),
+    "]"
+)
+
+# Plot results
+using Plots
+
+p1 = plot(
+    xs / pi,
+    ys,
+    V_predict,
+    linetype =
+    :contourf,
+    title = "V",
+    xlabel = "θ/π",
+    ylabel = "ω",
+    c = :bone_1
+);
+p1 = scatter!([-2 * pi, 0, 2 * pi] / pi, [0, 0, 0],
+    label = "Downward Equilibria", color = :green, markershape = :+);
+p1 = scatter!(
+    [-pi, pi] / pi, [0, 0], label = "Upward Equilibria", color = :red, markershape = :x);
+p2 = plot(
+    xs / pi,
+    ys,
+    dVdt_predict,
+    linetype = :contourf,
+    title = "dV/dt",
+    xlabel = "θ/π",
+    ylabel = "ω",
+    c = :binary
+);
+p2 = scatter!([-2 * pi, 0, 2 * pi] / pi, [0, 0, 0],
+    label = "Downward Equilibria", color = :green, markershape = :+);
+p2 = scatter!([-pi, pi] / pi, [0, 0], label = "Upward Equilibria",
+    color = :red, markershape = :x, legend = false);
+p3 = plot(
+    xs / pi,
+    ys,
+    dVdt_predict .< 0,
+    linetype = :contourf,
+    title = "dV/dt < 0",
+    xlabel = "θ/π",
+    ylabel = "ω",
+    colorbar = false,
+    linewidth = 0
+);
+p3 = scatter!([-2 * pi, 0, 2 * pi] / pi, [0, 0, 0],
+    label = "Downward Equilibria", color = :green, markershape = :+);
+p3 = scatter!([-pi, pi] / pi, [0, 0], label = "Upward Equilibria",
+    color = :red, markershape = :x, legend = false);
+plot(p1, p2, p3)
+=#
diff --git a/test/inverted_pendulum_ODESystem.jl b/test/inverted_pendulum_ODESystem.jl
new file mode 100644
index 0000000..5738c82
--- /dev/null
+++ b/test/inverted_pendulum_ODESystem.jl
@@ -0,0 +1,259 @@
+using LinearAlgebra
+using NeuralPDE, Lux, ModelingToolkit
+using Optimization, OptimizationOptimisers, OptimizationOptimJL, NLopt
+using NeuralLyapunov
+using Random
+using Test
+
+Random.seed!(200)
+
+println("Inverted Pendulum - Policy Search 2")
+
+######################### Define dynamics and domain ##########################
+
+@parameters ζ ω_0
+defaults = Dict([ζ => 0.5, ω_0 => 1.0])
+
+@variables t θ(t) τ(t) [input = true]
+Dt = Differential(t)
+DDt = Dt^2
+
+eqs = [DDt(θ) + 2ζ * Dt(θ) + ω_0^2 * sin(θ) ~ τ]
+
+@named driven_pendulum = ODESystem(
+    eqs,
+    t,
+    [θ, τ],
+    [ζ, ω_0];
+    defaults = defaults
+)
+
+bounds = [
+    θ ∈ (0, 2π),
+    Dt(θ) ∈ (-10.0, 10.0)
+]
+
+(open_loop_pendulum_dynamics, _), state_order, p_order = ModelingToolkit.generate_control_function(
+    driven_pendulum; simplify = true)
+
+upright_equilibrium = [π, 0.0]
+p = [defaults[param] for param in p_order]
+
+####################### Specify neural Lyapunov problem #######################
+
+# Define neural network discretization
+# We use an input layer that is periodic with period 2π with respect to θ
+dim_state = length(bounds)
+dim_hidden = 15
+dim_phi = 2
+dim_u = 1
+dim_output = dim_phi + dim_u
+chain = [Lux.Chain(
+             Lux.WrappedFunction(x -> vcat(
+                 transpose(sin.(x[1, :])),
+                 transpose(cos.(x[1, :])),
+                 transpose(x[2, :])
+             )),
+             Dense(3, dim_hidden, tanh),
+             Dense(dim_hidden, dim_hidden, tanh),
+             Dense(dim_hidden, 1, use_bias = false)
+         ) for _ in 1:dim_output]
+
+# Define neural network discretization
+strategy = GridTraining(0.1)
+discretization = PhysicsInformedNN(chain, strategy)
+
+# Define neural Lyapunov structure
+structure = PositiveSemiDefiniteStructure(
+    dim_phi;
+    pos_def = function (state, fixed_point)
+        θ, ω = state
+        θ_eq, ω_eq = fixed_point
+        log(1.0 + (sin(θ) - sin(θ_eq))^2 + (cos(θ) - cos(θ_eq))^2 + (ω - ω_eq)^2)
+    end
+)
+structure = add_policy_search(
+    structure,
+    dim_u
+)
+minimization_condition = DontCheckNonnegativity(check_fixed_point = false)
+
+# Define Lyapunov decrease condition
+decrease_condition = AsymptoticDecrease(strict = true)
+
+# Construct neural Lyapunov specification
+spec = NeuralLyapunovSpecification(
+    structure,
+    minimization_condition,
+    decrease_condition
+)
+
+############################# Construct PDESystem #############################
+
+pde_system, network_func = NeuralLyapunovPDESystem(
+    driven_pendulum,
+    bounds,
+    spec;
+    fixed_point = upright_equilibrium
+)
+
+######################## Construct OptimizationProblem ########################
+
+sym_prob = symbolic_discretize(pde_system, discretization)
+prob = discretize(pde_system, discretization)
+
+########################## Solve OptimizationProblem ##########################
+
+res = Optimization.solve(prob, OptimizationOptimisers.Adam(); maxiters = 300)
+prob = Optimization.remake(prob, u0 = res.u)
+res = Optimization.solve(prob, BFGS(); maxiters = 300)
+
+###################### Get numerical numerical functions ######################
+
+V_func, V̇_func, ∇V_func = NumericalNeuralLyapunovFunctions(
+    discretization.phi,
+    res.u,
+    network_func,
+    structure,
+    open_loop_pendulum_dynamics,
+    upright_equilibrium;
+    p = p
+)
+
+u = get_policy(discretization.phi, res.u, network_func, dim_u)
+
+################################## Simulate ###################################
+
+lb = [0.0, -10.0];
+ub = [2π, 10.0];
+xs = (-2π):0.1:(2π)
+ys = lb[2]:0.1:ub[2]
+states = Iterators.map(collect, Iterators.product(xs, ys))
+V_predict = vec(V_func(hcat(states...)))
+dVdt_predict = vec(V̇_func(hcat(states...)))
+
+#################################### Tests ####################################
+
+# Network structure should enforce positive definiteness
+@test V_func(upright_equilibrium) == 0.0
+@test min(V_func(upright_equilibrium), minimum(V_predict)) ≥ 0.0
+
+# Network structure should enforce periodicity in θ
+x0 = (ub .- lb) .* rand(2, 100) .+ lb
+@test all(isapprox.(V_func(x0), V_func(x0 .+ [2π, 0.0]); rtol = 1e-3))
+
+# Training should result in a fixed point at the upright equilibrium
+@test all(isapprox.(
+    open_loop_pendulum_dynamics(upright_equilibrium, u(upright_equilibrium), p, 0.0),
+    0.0; atol = 1e-8))
+@test V̇_func(upright_equilibrium) == 0.0
+
+# V̇ should be negative almost everywhere
+@test sum(dVdt_predict .> 0) / length(dVdt_predict) < 1.04e-3
+
+################################## Simulate ###################################
+
+using DifferentialEquations
+
+state_order = map(st -> istree(st) ? operation(st) : st, state_order)
+state_syms = Symbol.(state_order)
+
+closed_loop_dynamics = ODEFunction(
+    (x, p, t) -> open_loop_pendulum_dynamics(x, u(x), p, t);
+    sys = SciMLBase.SymbolCache(state_syms, Symbol.(p_order))
+)
+
+# Starting still at bottom
+downward_equilibrium = zeros(2)
+ode_prob = ODEProblem(closed_loop_dynamics, downward_equilibrium, [0.0, 70.0], p)
+sol = solve(ode_prob, Tsit5())
+# plot(sol)
+
+# Should make it to the top
+θ_end, ω_end = sol.u[end]
+x_end, y_end = sin(θ_end), -cos(θ_end)
+@test all(isapprox.([x_end, y_end, ω_end], [0.0, 1.0, 0.0]; atol = 1e-3))
+
+# Starting at a random point
+x0 = lb .+ rand(2) .* (ub .- lb)
+ode_prob = ODEProblem(closed_loop_dynamics, x0, [0.0, 100.0], p)
+sol = solve(ode_prob, Tsit5())
+# plot(sol)
+
+# Should make it to the top
+θ_end, ω_end = sol.u[end]
+x_end, y_end = sin(θ_end), -cos(θ_end)
+@test all(isapprox.([x_end, y_end, ω_end], [0.0, 1.0, 0.0]; atol = 1e-3))
+
+#=
+# Print statistics
+println("V(π, 0) = ", V_func(upright_equilibrium))
+println(
+    "f([π, 0], u([π, 0])) = ",
+    open_loop_pendulum_dynamics(upright_equilibrium, u(upright_equilibrium), p, 0.0)
+)
+println(
+    "V ∋ [",
+    min(V_func(upright_equilibrium),
+    minimum(V_predict)),
+    ", ",
+    maximum(V_predict),
+    "]"
+)
+println(
+    "V̇ ∋ [",
+    minimum(dVdt_predict),
+    ", ",
+    max(V̇_func(upright_equilibrium), maximum(dVdt_predict)),
+    "]"
+)
+
+# Plot results
+using Plots
+
+p1 = plot(
+    xs / pi,
+    ys,
+    V_predict,
+    linetype =
+    :contourf,
+    title = "V",
+    xlabel = "θ/π",
+    ylabel = "ω",
+    c = :bone_1
+);
+p1 = scatter!([-2 * pi, 0, 2 * pi] / pi, [0, 0, 0],
+    label = "Downward Equilibria", color = :green, markershape = :+);
+p1 = scatter!(
+    [-pi, pi] / pi, [0, 0], label = "Upward Equilibria", color = :red, markershape = :x);
+p2 = plot(
+    xs / pi,
+    ys,
+    dVdt_predict,
+    linetype = :contourf,
+    title = "dV/dt",
+    xlabel = "θ/π",
+    ylabel = "ω",
+    c = :binary
+);
+p2 = scatter!([-2 * pi, 0, 2 * pi] / pi, [0, 0, 0],
+    label = "Downward Equilibria", color = :green, markershape = :+);
+p2 = scatter!([-pi, pi] / pi, [0, 0], label = "Upward Equilibria",
+    color = :red, markershape = :x, legend = false);
+p3 = plot(
+    xs / pi,
+    ys,
+    dVdt_predict .< 0,
+    linetype = :contourf,
+    title = "dV/dt < 0",
+    xlabel = "θ/π",
+    ylabel = "ω",
+    colorbar = false,
+    linewidth = 0
+);
+p3 = scatter!([-2 * pi, 0, 2 * pi] / pi, [0, 0, 0],
+    label = "Downward Equilibria", color = :green, markershape = :+);
+p3 = scatter!([-pi, pi] / pi, [0, 0], label = "Upward Equilibria",
+    color = :red, markershape = :x, legend = false);
+plot(p1, p2, p3)
+=#
diff --git a/test/roa_estimation.jl b/test/roa_estimation.jl
index 43fb96e..8c341d4 100644
--- a/test/roa_estimation.jl
+++ b/test/roa_estimation.jl
@@ -70,8 +70,8 @@ V_func, V̇_func, ∇V_func = NumericalNeuralLyapunovFunctions(
     discretization.phi,
     res.u,
     network_func,
-    structure.V,
-    ODEFunction(f),
+    structure,
+    f,
     zeros(length(lb))
 )
 
diff --git a/test/runtests.jl b/test/runtests.jl
index 4887b76..ef1120d 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -7,4 +7,13 @@ using SafeTestsets
     @time @safetestset "Damped pendulum" begin
         include("damped_pendulum.jl")
     end
+    @time @safetestset "Region of attraction estimation" begin
+        include("roa_estimation.jl")
+    end
+    @time @safetestset "Policy search - inverted pendulum" begin
+        include("inverted_pendulum.jl")
+    end
+    @time @safetestset "Policy search - inverted pendulum 2" begin
+        include("inverted_pendulum_ODESystem.jl")
+    end
 end