Added cart-pole example

docs/dymos_book/_toc.yml

@@ -38,6 +38,7 @@ parts:
     - file: examples/ssto_moon_polynomial_controls/ssto_moon_polynomial_controls
     - file: examples/robertson_problem/robertson_problem
     - file: examples/water_rocket/water_rocket
+    - file: examples/cart_pole/cart_pole
 - caption: Feature Reference
   chapters:
   - file: features/phases/phases_list

docs/dymos_book/references.bib

@@ -575,3 +575,16 @@ @article{Berrut2004BarycentricLI
   year={2004},
   publisher={SIAM}
 }
+
+@article{Kelly2017,
+    title    = {{An introduction to trajectory optimization: How to do your own direct collocation}},
+    year     = {2017},
+    journal  = {SIAM Review},
+    author   = {Kelly, Matthew},
+    number   = {4},
+    pages    = {849--904},
+    volume   = {59},
+    doi      = {10.1137/16M1062569},
+    issn     = {00361445},
+    keywords = {Direct collocation, Direct transcription, Optimal control, Robotics, Trajectory optimization, Tutorial}
+}

dymos/examples/cart_pole/animate_cartpole.py

@@ -0,0 +1,74 @@
+import numpy as np
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+from matplotlib import animation
+
+
+def animate_cartpole(x, theta, force, interval=20, force_scaler=0.1, save_gif=False, gif_fps=20):
+    # x: time history of cart location, 1d vector
+    # theta: time history of pole angle, 1d vector
+    # force: control input force
+
+    # keep showing the final state for 100 ms
+    extend = 500 // interval
+    x = np.concatenate((x, np.ones(extend) * x[-1]))
+    theta = np.concatenate((theta, np.ones(extend) * theta[-1]))
+    force = np.concatenate((force, np.ones(extend) * force[-1]))
+
+    # cart parameters
+    cart_width = 0.25
+    cart_height = 0.15
+    lpole = 0.5
+
+    # path of the pole tip
+    x_pole = lpole * np.sin(theta) + x
+    y_pole = -lpole * np.cos(theta) + cart_height / 2.0
+
+    # force vector
+    arrow_direc = force * force_scaler
+    arrow_origin = x - arrow_direc
+    arrow_y = cart_height / 2
+
+    # x_lim and y_lim for figure
+    xlim = [min(min(x), min(x_pole)) - cart_width / 2 - 0.1, max(max(x), max(x_pole)) + cart_width / 2 + 0.1]
+    ylim = [cart_height / 2 - lpole - 0.05, cart_height / 2 + lpole + 0.05]
+
+    fig, ax = plt.subplots()
+
+    def init():
+        # initialize figure
+        ax.set_xlabel("x (m)")
+        ax.set_ylabel("y (m)")
+        ax.axis("equal")
+
+    def animate(i):
+        ax.clear()
+        # plot "floor"
+        ax.plot(xlim, [0.0, 0.0], color="k")
+        ax.set_xlim(xlim)
+        ax.set_ylim(ylim)
+        # plot cart
+        cart = ax.add_patch(
+            patches.Rectangle((-cart_width / 2.0 + x[i], 0.0), cart_width, cart_height, edgecolor="C0", linewidth=1, fill=False)
+        )
+        # plot pole
+        pole = ax.plot([x[i], x_pole[i]], [cart_height / 2, y_pole[i]], "o-", lw=2, color="C0")
+        # plot pole tip path
+        path = ax.plot(x_pole[:i], y_pole[:i], "--", lw=1, color="black")
+        # plot force vector
+        f = ax.arrow(
+            arrow_origin[i],
+            arrow_y,
+            arrow_direc[i],
+            0.0,
+            color="C1",
+            head_width=0.02,
+            head_length=0.03,
+            length_includes_head=True,
+        )
+        return pole, cart, path, f
+
+    anim = animation.FuncAnimation(fig, animate, init_func=init, frames=len(x), interval=interval, repeat=True)
+    if save_gif:
+        anim.save("cartpole.gif", dpi=300, writer=animation.PillowWriter(fps=gif_fps))
+    plt.show()

dymos/examples/cart_pole/cartpole_dynamics.py

@@ -0,0 +1,144 @@
+"""
+Cart-pole dynamics (ODE)
+"""
+
+import numpy as np
+import openmdao.api as om
+
+
+class CartPoleDynamics(om.ExplicitComponent):
+    """
+    Computes the time derivatives of states given state variables and control inputs.
+
+    Parameters
+    ----------
+    m_cart : float
+        Mass of cart.
+    m_pole : float
+        Mass of pole.
+    l_pole : float
+        Length of pole.
+    theta : 1d array
+        Angle of pole, 0 for vertical downward, positive counter clockwise.
+    theta_dot : 1d array
+        Angluar velocity of pole.
+    f : 1d array
+        x-wise force applied to the cart.
+
+    Returns
+    -------
+    x_dotdot : 1d array
+        Acceleration of cart in x direction.
+    theta_dotdot : 1d array
+        Angular acceleration of pole.
+    e_dot : 1d array
+        Rate of "energy" state.
+    """
+
+    def initialize(self):
+        self.options.declare(
+            "num_nodes", default=1, desc="number of nodes to be evaluated (i.e., length of vectors x, theta, etc)"
+        )
+        self.options.declare("g", default=9.81, desc="gravity constant")
+
+    def setup(self):
+        nn = self.options["num_nodes"]
+
+        # --- inputs ---
+        # cart-pole parameters
+        self.add_input("m_cart", shape=(1,), units="kg", desc="cart mass")
+        self.add_input("m_pole", shape=(1,), units="kg", desc="pole mass")
+        self.add_input("l_pole", shape=(1,), units="m", desc="pole length")
+        # state variables. States x, x_dot, energy have no influence on the outputs, so we don't need them as inputs.
+        self.add_input("theta", shape=(nn,), units="rad", desc="pole angle")
+        self.add_input("theta_dot", shape=(nn,), units="rad/s", desc="pole angle velocity")
+        # control input
+        self.add_input("f", shape=(nn,), units="N", desc="force applied to cart in x direction")
+
+        # --- outputs ---
+        # rate of states (accelerations)
+        self.add_output("x_dotdot", shape=(nn,), units="m/s**2", desc="x acceleration of cart")
+        self.add_output("theta_dotdot", shape=(nn,), units="rad/s**2", desc="angular acceleration of pole")
+        # also computes force**2, which will be integrated to compute the objective
+        self.add_output("e_dot", shape=(nn,), units="N**2", desc="square of force to be integrated")
+
+        # ---  partials ---.
+        # Jacobian of outputs w.r.t. state/control inputs is diagonal
+        # because each node (corresponds to time discretization) is independent
+        self.declare_partials(of=["*"], wrt=["theta", "theta_dot", "f"], method="exact", rows=np.arange(nn), cols=np.arange(nn))
+
+        # partials of outputs w.r.t. cart-pole parameters. We will use complex-step, but still declare the sparsity structure.
+        # NOTE: since the cart-pole parameters are fixed during optimization, these partials are not necessary to declare.
+        self.declare_partials(of=["*"], wrt=["m_cart", "m_pole", "l_pole"], method="cs", rows=np.arange(nn), cols=np.zeros(nn))
+        self.declare_coloring(wrt=["m_cart", "m_pole", "l_pole"], method="cs", show_summary=False)
+        self.set_check_partial_options(wrt=["m_cart", "m_pole", "l_pole"], method="fd", step=1e-7)
+
+    def compute(self, inputs, outputs):
+        g = self.options["g"]
+        mc = inputs["m_cart"]
+        mp = inputs["m_pole"]
+        lpole = inputs["l_pole"]
+        theta = inputs["theta"]
+        omega = inputs["theta_dot"]
+        f = inputs["f"]
+
+        sint = np.sin(theta)
+        cost = np.cos(theta)
+        det = mp * lpole * cost**2 - lpole * (mc + mp)
+        outputs["x_dotdot"] = (-mp * lpole * g * sint * cost - lpole * (f + mp * lpole * omega**2 * sint)) / det
+        outputs["theta_dotdot"] = ((mc + mp) * g * sint + cost * (f + mp * lpole * omega**2 * sint)) / det
+        outputs["e_dot"] = f**2
+
+    def compute_partials(self, inputs, jacobian):
+        g = self.options["g"]
+        mc = inputs["m_cart"]
+        mp = inputs["m_pole"]
+        lpole = inputs["l_pole"]
+        theta = inputs["theta"]
+        theta_dot = inputs["theta_dot"]
+        f = inputs["f"]
+
+        # --- derivatives of x_dotdot ---
+        # Collecting Theta Derivative
+        low = mp * lpole * np.cos(theta) ** 2 - lpole * mc - lpole * mp
+        dhigh = (
+            mp * g * lpole * np.sin(theta) ** 2 -
+            mp * g * lpole * np.cos(theta) ** 2 -
+            lpole**2 * mp * theta_dot**2 * np.cos(theta)
+        )
+        high = -mp * g * lpole * np.cos(theta) * np.sin(theta) - lpole * f - lpole**2 * mp * theta_dot**2 * np.sin(theta)
+        dlow = 2.0 * mp * lpole * np.cos(theta) * (-np.sin(theta))
+
+        jacobian["x_dotdot", "theta"] = (low * dhigh - high * dlow) / low**2
+        jacobian["x_dotdot", "theta_dot"] = (
+            -2.0 * theta_dot * lpole**2 * mp * np.sin(theta) / (mp * lpole * np.cos(theta) ** 2 - lpole * mc - lpole * mp)
+        )
+        jacobian["x_dotdot", "f"] = -lpole / (mp * lpole * np.cos(theta) ** 2 - lpole * mc - lpole * mp)
+
+        # --- derivatives of theta_dotdot ---
+        # Collecting Theta Derivative
+        low = mp * lpole * np.cos(theta) ** 2 - lpole * mc - lpole * mp
+        dlow = 2.0 * mp * lpole * np.cos(theta) * (-np.sin(theta))
+        high = (mc + mp) * g * np.sin(theta) + f * np.cos(theta) + mp * lpole * theta_dot**2 * np.sin(theta) * np.cos(theta)
+        dhigh = (
+            (mc + mp) * g * np.cos(theta) -
+            f * np.sin(theta) +
+            mp * lpole * theta_dot**2 * (np.cos(theta) ** 2 - np.sin(theta) ** 2)
+        )
+
+        jacobian["theta_dotdot", "theta"] = (low * dhigh - high * dlow) / low**2
+        jacobian["theta_dotdot", "theta_dot"] = (
+            2.0 *
+            theta_dot *
+            mp *
+            lpole *
+            np.sin(theta) *
+            np.cos(theta) /
+            (mp * lpole * np.cos(theta) ** 2 - lpole * mc - lpole * mp)
+        )
+        jacobian["theta_dotdot", "f"] = np.cos(theta) / (mp * lpole * np.cos(theta) ** 2 - lpole * mc - lpole * mp)
+
+        # --- derivatives of e_dot ---
+        jacobian["e_dot", "theta"] = 0.0
+        jacobian["e_dot", "theta_dot"] = 0.0
+        jacobian["e_dot", "f"] = 2.0 * f

dymos/examples/cart_pole/test/test_cartpole_dynamics.py

@@ -0,0 +1,39 @@
+import unittest
+
+import openmdao.api as om
+from openmdao.utils.assert_utils import assert_near_equal, assert_check_partials
+from openmdao.utils.testing_utils import use_tempdirs
+
+from dymos.examples.cart_pole.cartpole_dynamics import CartPoleDynamics
+
+
+@use_tempdirs
+class TestCartPoleDynamics(unittest.TestCase):
+    def test_cartpole_ode(self):
+
+        p = om.Problem()
+        p.model.add_subsystem("dynamics", CartPoleDynamics(num_nodes=2), promotes=["*"])
+        # set input values
+        p.model.set_input_defaults("m_cart", 1.0, units="kg")
+        p.model.set_input_defaults("m_pole", 0.3, units="kg")
+        p.model.set_input_defaults("l_pole", 0.5, units="m")
+        p.model.set_input_defaults("theta", [0.0, 1.0], units="rad")
+        p.model.set_input_defaults("theta_dot", [0.0, 0.1], units="rad/s")
+        p.model.set_input_defaults("f", [10, -10], units="N")
+
+        # run model
+        p.setup(check=False)
+        p.run_model()
+        # get outputs
+        x_dotdot = p.get_val("x_dotdot", units="m/s**2")
+        theta_dotdot = p.get_val("theta_dotdot", units="rad/s**2")
+        assert_near_equal(x_dotdot, [10.0, -7.14331021], tolerance=1e-6)
+        assert_near_equal(theta_dotdot, [-20.0, -8.79056677], tolerance=1e-6)
+
+        # check partials
+        partials = p.check_partials(compact_print=True, method="cs")
+        assert_check_partials(partials, atol=1e-5, rtol=1e-5)
+
+
+if __name__ == "___main__":
+    unittest.main()

dymos/examples/cart_pole/test/test_cartpole_opt.py

@@ -0,0 +1,103 @@
+import unittest
+
+import numpy as np
+
+import openmdao.api as om
+from openmdao.utils.assert_utils import assert_near_equal
+from openmdao.utils.testing_utils import use_tempdirs, require_pyoptsparse
+
+import dymos as dm
+from dymos.examples.cart_pole.cartpole_dynamics import CartPoleDynamics
+
+
+@use_tempdirs
+class TestCartPoleOptimization(unittest.TestCase):
+    @require_pyoptsparse(optimizer="SNOPT")
+    def test_optimization(self):
+
+        p = om.Problem()
+
+        # --- instantiate trajectory and phase, setup transcription ---
+        traj = dm.Trajectory()
+        p.model.add_subsystem("traj", traj)
+        phase = dm.Phase(
+            transcription=dm.GaussLobatto(num_segments=40, order=3, compressed=True, solve_segments=False),
+            ode_class=CartPoleDynamics,
+        )
+        # NOTE: set solve_segments=True to do solver-based shooting
+        traj.add_phase("phase", phase)
+
+        # --- set state and control variables ---
+        phase.set_time_options(fix_initial=True, fix_duration=True, duration_val=2.0, units="s")
+        # declare state variables. You can also set lower/upper bounds and scalings here.
+        phase.add_state("x", fix_initial=True, lower=-2, upper=2, rate_source="x_dot", shape=(1,), ref=1, defect_ref=1, units="m")
+        phase.add_state("x_dot", fix_initial=True, rate_source="x_dotdot", shape=(1,), ref=1, defect_ref=1, units="m/s")
+        phase.add_state("theta", fix_initial=True, rate_source="theta_dot", shape=(1,), ref=1, defect_ref=1, units="rad")
+        phase.add_state("theta_dot", fix_initial=True, rate_source="theta_dotdot", shape=(1,), ref=1, defect_ref=1, units="rad/s")
+        phase.add_state(
+            "energy", fix_initial=True, rate_source="e_dot", shape=(1,), ref=1, defect_ref=1, units="N**2*s"
+        )  # integration of force**2. This does not have the energy unit, but I call it "energy" anyway.
+
+        # declare control inputs
+        phase.add_control("f", fix_initial=False, rate_continuity=False, lower=-20, upper=20, shape=(1,), ref=10, units="N")
+
+        # add cart-pole parameters (set static_target=True because these params are not time-depencent)
+        phase.add_parameter("m_cart", val=1.0, units="kg", static_target=True)
+        phase.add_parameter("m_pole", val=0.3, units="kg", static_target=True)
+        phase.add_parameter("l_pole", val=0.5, units="m", static_target=True)
+
+        # --- set terminal constraint ---
+        # alternatively, you can impose those by setting `fix_final=True` in phase.add_state()
+        phase.add_boundary_constraint("x", loc="final", equals=1, ref=1.0, units="m")  # final horizontal displacement
+        phase.add_boundary_constraint("theta", loc="final", equals=np.pi, ref=1.0, units="rad")  # final pole angle
+        phase.add_boundary_constraint("x_dot", loc="final", equals=0, ref=1.0, units="m/s")  # 0 velocity at the and
+        phase.add_boundary_constraint("theta_dot", loc="final", equals=0, ref=0.1, units="rad/s")  # 0 angular velocity at the end
+        phase.add_boundary_constraint("f", loc="final", equals=0, ref=1.0, units="N")  # 0 force at the end
+
+        # --- set objective function ---
+        # we minimize the integral of force**2.
+        phase.add_objective("energy", loc="final", ref=100)
+
+        # --- configure optimizer ---
+        p.driver = om.pyOptSparseDriver()
+        p.driver.options["optimizer"] = "IPOPT"
+        # IPOPT options
+        p.driver.opt_settings['mu_init'] = 1e-1
+        p.driver.opt_settings['max_iter'] = 600
+        p.driver.opt_settings['constr_viol_tol'] = 1e-6
+        p.driver.opt_settings['compl_inf_tol'] = 1e-6
+        p.driver.opt_settings['tol'] = 1e-5
+        p.driver.opt_settings['print_level'] = 0
+        p.driver.opt_settings['nlp_scaling_method'] = 'gradient-based'
+        p.driver.opt_settings['alpha_for_y'] = 'safer-min-dual-infeas'
+        p.driver.opt_settings['mu_strategy'] = 'monotone'
+        p.driver.opt_settings['bound_mult_init_method'] = 'mu-based'
+        p.driver.options['print_results'] = False
+
+        # declare total derivative coloring to accelerate the UDE linear solves
+        p.driver.declare_coloring()
+
+        p.setup(check=False)
+
+        # --- set initial guess ---
+        # The initial condition of cart-pole (i.e., state values at time 0) is set here
+        # because we set `fix_initial=True` when declaring the states.
+        p.set_val("traj.phase.t_initial", 0.0)  # set initial time to 0.
+        p.set_val("traj.phase.states:x", phase.interp(xs=[0, 1, 2], ys=[0, 1, 1], nodes="state_input"), units="m")
+        p.set_val("traj.phase.states:x_dot", phase.interp(xs=[0, 1, 2], ys=[0, 0.1, 0], nodes="state_input"), units="m/s")
+        p.set_val("traj.phase.states:theta", phase.interp(xs=[0, 1, 2], ys=[0, np.pi/2, np.pi], nodes="state_input"), units="rad")
+        p.set_val("traj.phase.states:theta_dot", phase.interp(xs=[0, 1, 2], ys=[0, 1, 0], nodes="state_input"), units="rad/s")
+        p.set_val("traj.phase.states:energy", phase.interp(xs=[0, 1, 2], ys=[0, 30, 60], nodes="state_input"))
+        p.set_val("traj.phase.controls:f", phase.interp(xs=[0, 1, 2], ys=[3, -1, 0], nodes="control_input"), units="N")
+
+        # --- run optimization ---
+        dm.run_problem(p, run_driver=True, simulate=False, simulate_kwargs={"method": "Radau", "times_per_seg": 10})
+
+        # --- check outputs ---
+        # objective value
+        obj = p.get_val("traj.phase.states:energy", units="N**2*s")[-1]
+        assert_near_equal(obj, 58.8839489745, tolerance=1e-3)
+
+
+if __name__ == "___main__":
+    unittest.main()