.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "dynamics/cartpole_tutorial.py" .. LINE NUMBERS ARE GIVEN BELOW. .. only:: html .. note:: :class: sphx-glr-download-link-note Click :ref:`here ` to download the full example code .. rst-class:: sphx-glr-example-title .. _sphx_glr_dynamics_cartpole_tutorial.py: Cartpole Tutorial ================== .. GENERATED FROM PYTHON SOURCE LINES 6-11 .. code-block:: default import torch, pypose as pp import math, matplotlib.pyplot as plt device = torch.device("cuda" if torch.cuda.is_available() else "cpu") .. GENERATED FROM PYTHON SOURCE LINES 12-16 Preparation ------------- Create class for cart-pole dynamics .. GENERATED FROM PYTHON SOURCE LINES 16-57 .. code-block:: default class CartPole(pp.module.NLS): def __init__(self, dt, length, cartmass, polemass, gravity): super().__init__() self.tau = dt self.length = length self.cartmass = cartmass self.polemass = polemass self.gravity = gravity self.polemassLength = self.polemass * self.length self.totalMass = self.cartmass + self.polemass def state_transition(self, state, input, t = None): x, xDot, theta, thetaDot = state force = input.squeeze() costheta = theta.cos() sintheta = theta.sin() temp = (force + self.polemassLength * thetaDot**2 * sintheta) / self.totalMass thetaAcc = (self.gravity * sintheta - costheta * temp) / \ (self.length * (4.0 / 3.0 - self.polemass * costheta**2 / self.totalMass)) xAcc = temp - self.polemassLength * thetaAcc * costheta / self.totalMass _dstate = torch.stack((xDot, xAcc, thetaDot, thetaAcc)) return state + _dstate * self.tau def observation(self, state, input, t = None): return state def subPlot(ax, x, y, xlabel=None, ylabel=None): x = x.detach().cpu().numpy() y = y.detach().cpu().numpy() ax.plot(x, y) ax.set_xlabel(xlabel) ax.set_ylabel(ylabel) .. GENERATED FROM PYTHON SOURCE LINES 58-60 Create parameters for cart pole trajectory -------------------------------------------- .. GENERATED FROM PYTHON SOURCE LINES 60-70 .. code-block:: default dt = 0.01 # Delta t len = 1.5 # Length of pole m_cart = 20 # Mass of cart m_pole = 10 # Mass of pole g = 9.81 # Accerleration due to gravity N = 1000 # Number of time steps .. GENERATED FROM PYTHON SOURCE LINES 71-72 Time and input .. GENERATED FROM PYTHON SOURCE LINES 72-77 .. code-block:: default time = torch.arange(0, N, device=device) * dt input = torch.sin(time) .. GENERATED FROM PYTHON SOURCE LINES 78-79 Initial state .. GENERATED FROM PYTHON SOURCE LINES 79-83 .. code-block:: default state = torch.zeros(N, 4, dtype=float, device=device) state[0] = torch.tensor([0, 0, math.pi, 0], dtype=float, device=device) .. GENERATED FROM PYTHON SOURCE LINES 84-85 Create dynamics solver object .. GENERATED FROM PYTHON SOURCE LINES 85-88 .. code-block:: default model = CartPole(dt, len, m_cart, m_pole, g).to(device) .. GENERATED FROM PYTHON SOURCE LINES 89-90 Calculate trajectory .. GENERATED FROM PYTHON SOURCE LINES 90-94 .. code-block:: default for i in range(N - 1): state[i + 1], _ = model(state[i], input[i]) .. GENERATED FROM PYTHON SOURCE LINES 95-96 Jacobian computation - Find jacobians at the last step .. GENERATED FROM PYTHON SOURCE LINES 96-101 .. code-block:: default model.set_refpoint(state=state[-1,:], input=input[-1], t=time[-1]) vars = ['A', 'B', 'C', 'D', 'c1', 'c2'] [print(v, getattr(model, v)) for v in vars] .. rst-class:: sphx-glr-script-out .. code-block:: none A tensor([[ 1.0000e+00, 1.0000e-02, 0.0000e+00, 0.0000e+00], [ 0.0000e+00, 1.0000e+00, -3.2700e-02, -3.4183e-07], [ 0.0000e+00, 0.0000e+00, 1.0000e+00, 1.0000e-02], [ 0.0000e+00, 0.0000e+00, -6.5400e-02, 1.0000e+00]], device='cuda:0', dtype=torch.float64) B tensor([[0.0000], [0.0004], [0.0000], [0.0002]], device='cuda:0') C tensor([[1., 0., 0., 0.], [0., 1., 0., 0.], [0., 0., 1., 0.], [0., 0., 0., 1.]], device='cuda:0', dtype=torch.float64) D tensor([[0.], [0.], [0.], [0.]], device='cuda:0') c1 tensor([0.0000, 0.1027, 0.0000, 0.2055], device='cuda:0', dtype=torch.float64) c2 tensor([0., 0., 0., 0.], device='cuda:0', dtype=torch.float64) [None, None, None, None, None, None] .. GENERATED FROM PYTHON SOURCE LINES 102-103 Create time plots to show dynamics .. GENERATED FROM PYTHON SOURCE LINES 103-110 .. code-block:: default f, ax = plt.subplots(nrows=4, sharex=True) x, xdot, theta, thetadot = state.T subPlot(ax[0], time, x, ylabel='X') subPlot(ax[1], time, xdot, ylabel='X dot') subPlot(ax[2], time, theta, ylabel='Theta') subPlot(ax[3], time, thetadot, ylabel='Theta dot', xlabel='Time') plt.show() .. image-sg:: /dynamics/images/sphx_glr_cartpole_tutorial_001.png :alt: cartpole tutorial :srcset: /dynamics/images/sphx_glr_cartpole_tutorial_001.png :class: sphx-glr-single-img .. rst-class:: sphx-glr-timing **Total running time of the script:** ( 0 minutes 0.170 seconds) .. _sphx_glr_download_dynamics_cartpole_tutorial.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: cartpole_tutorial.py ` .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: cartpole_tutorial.ipynb ` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_