{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# **Cart-Pole System and Implementation in Crocoddyl**\n",
    "\n",
    "This notebook implements the cartpole system using **Crocoddyl**, a differential dynamic programming solver for robotic control. The key features of this notebook are\n",
    " - **System Modeling**: Define the cartpole dynamics as an action model in Crocoddyl.\n",
    " - **Cost Function**: Implement a cost function that penalizes deviations from the upright position and large control inputs within an action model.\n",
    " - **Derivatives Computation**: Compute analytical derivatives of the dynamics and cost function for efficient optimization.\n",
    " - **Numerical differentiation**: Validate analytical derivatives using Crocoddyl's `ActionModelNumDiff` class.\n",
    " - **Simulation & Optimization**: Simulate the system and solves an optimal control problem using Crocoddyl.\n",
    "\n",
    "## Overview of the Cart-Pole System\n",
    "\n",
    "The cart-pole system consists of a cart that can move along a horizontal track and a pole attached to the cart via a pivot joint. The goal is to balance the pole in the upright position by applying forces to the cart.\n",
    "\n",
    "<div style=\"text-align:center;\">\n",
    "    <img src=\"https://www.ashwinnarayan.com/img/cart-pole.png\" width=\"300\" height=\"300\">\n",
    "</div>\n",
    "\n",
    "## Equations of Motion\n",
    "\n",
    "The state of the system is given by:\n",
    "\n",
    "- $y$: Position of the cart\n",
    "- $\\theta$: Angle of the pole from the upright position\n",
    "- $\\dot{y}$: Velocity of the cart\n",
    "- $\\dot{\\theta}$: Angular velocity of the pole\n",
    "\n",
    "The equations of motion are derived from Newtonian mechanics:\n",
    "\n",
    "\\begin{aligned}\n",
    "\\mu &= m_{\\text{cart}} + m_{\\text{pole}} \\sin^2(\\theta) \\\\\n",
    "\\ddot{y} &= \\frac{f}{\\mu} + \\frac{m_{\\text{pole}} g \\cos(\\theta) \\sin(\\theta)}{\\mu} - \\frac{m_{\\text{pole}} l_{\\text{pole}} \\sin(\\theta) \\dot{\\theta}^2}{\\mu} \\\\\n",
    "\\ddot{\\theta} &= \\frac{f \\cos(\\theta)}{l_{\\text{pole}} \\mu} + \\frac{(m_{\\text{cart}} + m_{\\text{pole}}) g \\sin(\\theta)}{l_{\\text{pole}} \\mu} - \\frac{m_{\\text{pole}} \\cos(\\theta) \\sin(\\theta) \\dot{\\theta}^2}{\\mu}\n",
    "\\end{aligned}\n",
    "\n",
    "where:\n",
    " - $\\ddot{y}$, $\\ddot{\\theta}$ are accelerations\n",
    " - $f$ is the control force applied to the cart\n",
    " - $m_{\\text{cart}}$ is the mass of the cart\n",
    " - $m_{\\text{pole}}$ is the mass of the pole\n",
    " - $l_{\\text{pole}}$ is the length of the pole\n",
    " - $g$ is the gravitational acceleration\n",
    "\n",
    "The system dynamics update is given by:\n",
    "\n",
    "\\begin{aligned}\n",
    "y_{t+1} &= y_t + \\Delta t\\, \\dot{y}_{t+1} \\\\\n",
    "\\theta_{t+1} &= \\theta_t + \\Delta t\\, \\dot{\\theta}_{t+1} \\\\\n",
    "\\dot{y}_{t+1} &= \\dot{y}_t + \\Delta t\\, \\ddot{y} \\\\\n",
    "\\dot{\\theta}_{t+1} &= \\dot{\\theta}_t + \\Delta t\\, \\ddot{\\theta}\n",
    "\\end{aligned}\n",
    "\n",
    "### Cost Function\n",
    "A quadratic cost function is used to regulate the system:\n",
    "\n",
    "\\begin{equation}\n",
    "J = \\frac{1}{2} \\sum w_i r_i^2\n",
    "\\end{equation}\n",
    "\n",
    "where residuals $r$ are:\n",
    "\n",
    "\\begin{equation}\n",
    "r = \\begin{bmatrix} \\sin \\theta, 1 - \\cos \\theta, y, \\dot{y}, \\dot{\\theta}, f \\end{bmatrix}\n",
    "\\end{equation}\n",
    "\n",
    "and weights $w$ penalize deviations.\n",
    "\n",
    "\n",
    "## Implementation in Crocoddyl\n",
    "\n",
    "### ActionModelCartpole\n",
    "\n",
    "This class models the cart-pole system using Crocoddyl’s `ActionModelAbstract`. It defines the system dynamics and cost function.\n",
    "\n",
    "Attributes:\n",
    " - `Δt`: Time step for integration.\n",
    " - `m_cart`: Mass of the cart.\n",
    " - `m_pole`: Mass of the pole.\n",
    " - `l_pole`: Length of the pole.\n",
    " - `grav`: Gravitational acceleration.\n",
    " - `costWeights`: Weights for the cost function.\n",
    "\n",
    "Methods:\n",
    " - `calc(data, x, u)`: Computes the system dynamics and cost.\n",
    " - `calcDiff(data, x, u)`: Computes the derivatives of the system dynamics and cost function.\n",
    " - `createData()`: Creates an instance of `ActionDataCartpole`.\n",
    "\n",
    "Below, you can see"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import crocoddyl\n",
    "import numpy as np\n",
    "\n",
    "class ActionModelCartpole(crocoddyl.ActionModelAbstract):\n",
    "    \"\"\"\n",
    "    Defines the action model for the cart-pole system in Crocoddyl.\n",
    "    \n",
    "    This model simulates a cart moving along a horizontal track with a \n",
    "    freely rotating pole attached to it. The goal is to control the force \n",
    "    applied to the cart to stabilize the pole.\n",
    "    \n",
    "    State vector: [y, θ, ẏ, θ̇]\n",
    "    Control input: [f] (force applied to the cart)\n",
    "    \"\"\"\n",
    "    def __init__(self):\n",
    "        \"\"\"\n",
    "        Initializes the cart-pole action model with system parameters and cost weights.\n",
    "        \"\"\"\n",
    "        super().__init__(crocoddyl.StateVector(4), 1, 6)  # nu = 1; nr = 6\n",
    "\n",
    "        self.Δt = 5e-2\n",
    "        self.m_cart = 1.0\n",
    "        self.m_pole = 0.1\n",
    "        self.l_pole = 0.5\n",
    "        self.grav = 9.81\n",
    "        self.costWeights = [\n",
    "            1.0,\n",
    "            1.0,\n",
    "            0.1,\n",
    "            0.001,\n",
    "            0.001,\n",
    "            1.0,\n",
    "        ]  # sin, 1-cos, x, xdot, thdot, f\n",
    "\n",
    "    def calc(self, data, x, u=None):\n",
    "        \"\"\"\n",
    "        Computes the next state and the cost function.\n",
    "\n",
    "        Args:\n",
    "            data (ActionDataCartpole): Data structure to store intermediate results.\n",
    "            x (numpy.ndarray): State vector [y, θ, ẏ, θ̇].\n",
    "            u (numpy.ndarray, optional): Control input [f]. Defaults to None.\n",
    "\n",
    "        Updates:\n",
    "            - Next state xnext.\n",
    "            - Cost residuals r.\n",
    "            - Total cost value.\n",
    "        \"\"\"\n",
    "        if u is not None:\n",
    "            # Getting the state and control variables\n",
    "            y, θ, ẏ, θ̇, f = x[0], x[1], x[2], x[3], u[0]\n",
    "            # Shortname for system parameters\n",
    "            Δt, m_cart, m_pole, l_pole, grav, w = self.Δt, self.m_cart, self.m_pole, self.l_pole, self.grav, self.costWeights\n",
    "            sin_θ, cos_θ = np.sin(θ), np.cos(θ)\n",
    "            # Computing the cartpole dynamics\n",
    "            data.μ = m_cart + m_pole * sin_θ * sin_θ\n",
    "            data.ÿ̈ = f / data.μ\n",
    "            data.ÿ̈ += m_pole * grav * cos_θ * sin_θ / data.μ\n",
    "            data.ÿ̈ -= m_pole * l_pole * (sin_θ * θ̇ ** 2 / data.μ)\n",
    "            data.θ̈ = (f / l_pole) * cos_θ / data.μ\n",
    "            data.θ̈ += ((m_cart + m_pole) * grav / l_pole) * sin_θ / data.μ\n",
    "            data.θ̈ -= m_pole * cos_θ * sin_θ * θ̇**2 / data.μ\n",
    "            data.ẏ_next = ẏ + Δt * data.ÿ̈\n",
    "            data.θ̇_next = θ̇ + Δt * data.θ̈\n",
    "            data.y_next = y + Δt * data.ẏ_next\n",
    "            data.θ_next = θ + Δt * data.θ̇_next\n",
    "            data.xnext[:] = np.array([data.y_next, data.θ_next, data.ẏ_next, data.θ̇_next])\n",
    "            # Computing the cost residual and value\n",
    "            data.r[:] = w * np.array([sin_θ, 1.0 - cos_θ, y, ẏ, θ̇, f])\n",
    "            data.cost = 0.5 * sum(data.r ** 2)\n",
    "        else:\n",
    "            # Getting the state and control variables\n",
    "            y, θ, ẏ, θ̇ = x[0], x[1], x[2], x[3]\n",
    "            w = self.costWeights\n",
    "            sin_θ, cos_θ = np.sin(θ), np.cos(θ)\n",
    "            data.xnext[:] = x\n",
    "            # Computing the cost residual and value\n",
    "            data.r[:] = w * np.array([sin_θ, 1.0 - cos_θ, y, ẏ, θ̇, 0.0])\n",
    "            data.cost = 0.5 * sum(data.r ** 2)\n",
    "\n",
    "    def calcDiff(self, data, x, u=None):\n",
    "        \"\"\"\n",
    "        Computes the derivatives of the dynamics and the cost function.\n",
    "\n",
    "        Args:\n",
    "            data (ActionDataCartpole): Data structure to store intermediate results.\n",
    "            x (numpy.ndarray): State vector [y, θ, ẏ, θ̇].\n",
    "            u (numpy.ndarray, optional): Control input [f]. Defaults to None.\n",
    "\n",
    "        Updates:\n",
    "            - Derivetives of the dynamics.\n",
    "            - Derivatives of the cost function.\n",
    "        \"\"\"\n",
    "        if u is not None:\n",
    "            # Getting the state and control variables\n",
    "            y, θ, ẏ, θ̇, f = x[0], x[1], x[2], x[3], u[0]\n",
    "            # Shortname for system parameters\n",
    "            Δt, m_cart, m_pole, l_pole, grav, w = self.Δt, self.m_cart, self.m_pole, self.l_pole, self.grav, self.costWeights\n",
    "            sin_θ, cos_θ = np.sin(θ), np.cos(θ)\n",
    "            # Computing the derivative of the cartpole dynamics\n",
    "            data.dμ_dθ = 2.0 * m_pole * sin_θ * cos_θ\n",
    "            data.dÿ̈_dy = 0.0\n",
    "            data.dÿ̈_dθ = m_pole * grav * (cos_θ**2 - sin_θ**2) / data.μ\n",
    "            data.dÿ̈_dθ -= m_pole * l_pole * cos_θ * θ̇ * θ̇ / data.μ\n",
    "            data.dÿ̈_dθ -= data.dμ_dθ * data.ÿ̈ / data.μ\n",
    "            data.dÿ̈_dẏ = 0.0\n",
    "            data.dÿ̈_dθ̇ = -2.0 * m_pole * l_pole * sin_θ * θ̇ / data.μ\n",
    "            data.dÿ̈_du = 1.0 / data.μ\n",
    "            data.dθ̈_dy = 0.0\n",
    "            data.dθ̈_dθ = -(f / l_pole) * sin_θ / data.μ\n",
    "            data.dθ̈_dθ += ((m_cart + m_pole) * grav / l_pole) * cos_θ / data.μ\n",
    "            data.dθ̈_dθ -= m_pole * (cos_θ**2 - sin_θ**2) * θ̇ * θ̇ / data.μ\n",
    "            data.dθ̈_dθ -= data.dμ_dθ * data.θ̈ / data.μ\n",
    "            data.dθ̈_dẏ = 0.0\n",
    "            data.dθ̈_dθ̇ = -2.0 * m_pole * cos_θ * sin_θ * θ̇ / data.μ\n",
    "            data.dθ̈_du = cos_θ / (l_pole * data.μ)\n",
    "            data.dẏ_next_dy = Δt * data.dÿ̈_dy\n",
    "            data.dẏ_next_dθ = Δt * data.dÿ̈_dθ\n",
    "            data.dẏ_next_dẏ = 1.0 + Δt * data.dÿ̈_dẏ\n",
    "            data.dẏ_next_dθ̇ = Δt * data.dÿ̈_dθ̇\n",
    "            data.dẏ_next_du = Δt * data.dÿ̈_du\n",
    "            data.dθ̇_next_dy = Δt * data.dθ̈_dy \n",
    "            data.dθ̇_next_dθ = Δt * data.dθ̈_dθ\n",
    "            data.dθ̇_next_dẏ = Δt * data.dθ̈_dẏ\n",
    "            data.dθ̇_next_dθ̇ = 1.0 + Δt * data.dθ̈_dθ̇\n",
    "            data.dθ̇_next_du = Δt * data.dθ̈_du\n",
    "            data.dy_next_dy = 1.0 + Δt * data.dẏ_next_dy\n",
    "            data.dy_next_dθ = Δt * data.dẏ_next_dθ\n",
    "            data.dy_next_dẏ = Δt * data.dẏ_next_dẏ\n",
    "            data.dy_next_dθ̇ = Δt * data.dẏ_next_dθ̇\n",
    "            data.dy_next_du = Δt * data.dẏ_next_du\n",
    "            data.dθ_next_dy = Δt * data.dθ̇_next_dy\n",
    "            data.dθ_next_dθ = 1.0 + Δt * data.dθ̇_next_dθ\n",
    "            data.dθ_next_dẏ = Δt * data.dθ̇_next_dẏ\n",
    "            data.dθ_next_dθ̇ = Δt * data.dθ̇_next_dθ̇\n",
    "            data.dθ_next_du = Δt * data.dθ̇_next_du\n",
    "            data.Fx[:, :] = np.array([[data.dy_next_dy, data.dy_next_dθ, data.dy_next_dẏ, data.dy_next_dθ̇],\n",
    "                                      [data.dθ_next_dy, data.dθ_next_dθ, data.dθ_next_dẏ, data.dθ_next_dθ̇],\n",
    "                                      [data.dẏ_next_dy, data.dẏ_next_dθ, data.dẏ_next_dẏ, data.dẏ_next_dθ̇],\n",
    "                                      [data.dθ̇_next_dy, data.dθ̇_next_dθ, data.dθ̇_next_dẏ, data.dθ̇_next_dθ̇]])\n",
    "            data.Fu[:] = np.array([data.dy_next_du, data.dθ_next_du, data.dẏ_next_du, data.dθ̇_next_du])\n",
    "            # Computing derivatives of the cost function\n",
    "            w0_2, w1_2, w2_2, w3_2, w4_2, w5_2 = w[0] * w[0], w[1] * w[1], w[2] * w[2], w[3] * w[3], w[4] * w[4], w[5] * w[5]\n",
    "            data.Lx[0] = w2_2 * y\n",
    "            data.Lx[1] = w0_2 * sin_θ * cos_θ + w1_2 * (1.0 - cos_θ) * sin_θ\n",
    "            data.Lx[2] = w3_2 * ẏ\n",
    "            data.Lx[3] = w4_2 * θ̇\n",
    "            data.Lu[0] = w5_2 * f\n",
    "            data.Lxx[0, 0] = w2_2\n",
    "            data.Lxx[1, 1] = w0_2 * (cos_θ * cos_θ - sin_θ * sin_θ)\n",
    "            data.Lxx[1, 1] += w1_2 * ((1.0 - cos_θ) * cos_θ + sin_θ * sin_θ)\n",
    "            data.Lxx[2, 2] = w3_2\n",
    "            data.Lxx[3, 3] = w4_2\n",
    "            data.Luu[:] = w5_2\n",
    "        else:\n",
    "            # Getting the state and control variables\n",
    "            y, θ, ẏ, θ̇ = x[0], x[1], x[2], x[3]\n",
    "            w = self.costWeights\n",
    "            sin_θ, cos_θ = np.sin(θ), np.cos(θ)\n",
    "            # Computing the derivative of the cartpole dynamics\n",
    "            for i in range(self.state.ndx):\n",
    "                data.Fx[i, i] = 1.0\n",
    "            # Computing derivatives of the cost function\n",
    "            w0_2, w1_2, w2_2, w3_2, w4_2, w5_2 = w[0] * w[0], w[1] * w[1], w[2] * w[2], w[3] * w[3], w[4] * w[4], w[5] * w[5]\n",
    "            data.Lx[0] = w2_2 * y\n",
    "            data.Lx[1] = w0_2 * sin_θ * cos_θ + w1_2 * (1.0 - cos_θ) * sin_θ\n",
    "            data.Lx[2] = w3_2 * ẏ\n",
    "            data.Lx[3] = w4_2 * θ̇\n",
    "            data.Lxx[0, 0] = w2_2\n",
    "            data.Lxx[1, 1] = w0_2 * (cos_θ * cos_θ - sin_θ * sin_θ)\n",
    "            data.Lxx[1, 1] += w1_2 * ((1.0 - cos_θ) * cos_θ + sin_θ * sin_θ)\n",
    "            data.Lxx[2, 2] = w3_2\n",
    "            data.Lxx[3, 3] = w4_2\n",
    "\n",
    "    def createData(self):\n",
    "        \"\"\"\n",
    "        Creates the action data structure for the cart-pole model.\n",
    "\n",
    "        Returns:\n",
    "            ActionDataCartpole: Data structure to store intermediate computations.\n",
    "        \"\"\"\n",
    "        return ActionDataCartpole(self)\n",
    "\n",
    "\n",
    "class ActionDataCartpole(crocoddyl.ActionDataAbstract):\n",
    "    \"\"\"\n",
    "    Data structure for storing intermediate computations of the cart-pole dynamics.\n",
    "    \"\"\"\n",
    "    def __init__(self, model):\n",
    "        \"\"\"\n",
    "        Initializes the data structure with default values.\n",
    "\n",
    "        Args:\n",
    "            model (ActionModelCartpole): Cart-pole action model.\n",
    "        \"\"\"\n",
    "        super().__init__(model)\n",
    "        self.μ = 0.0\n",
    "        self.ÿ̈ = 0.0\n",
    "        self.θ̈ = 0.0\n",
    "        self.ẏ_next = 0.0\n",
    "        self.θ̇_next = 0.0\n",
    "        self.y_next = 0.0\n",
    "        self.θ_next = 0.0\n",
    "        self.dμ_dθ = 0.0\n",
    "        self.dÿ̈_dy = 0.0\n",
    "        self.dÿ̈_dθ = 0.0\n",
    "        self.dÿ̈_dθ = 0.0\n",
    "        self.dÿ̈_dθ = 0.0\n",
    "        self.dÿ̈_dẏ = 0.0\n",
    "        self.dÿ̈_dθ̇ = 0.0\n",
    "        self.dÿ̈_du = 0.0\n",
    "        self.dθ̈_dy = 0.0\n",
    "        self.dθ̈_dθ = 0.0\n",
    "        self.dθ̈_dθ = 0.0\n",
    "        self.dθ̈_dθ = 0.0\n",
    "        self.dθ̈_dθ = 0.0\n",
    "        self.dθ̈_dẏ = 0.0\n",
    "        self.dθ̈_dθ̇ = 0.0\n",
    "        self.dẏ_next_dy = 0.0\n",
    "        self.dẏ_next_dθ = 0.0\n",
    "        self.dẏ_next_dẏ = 0.0\n",
    "        self.dẏ_next_dθ̇ = 0.0\n",
    "        self.dẏ_next_du = 0.0\n",
    "        self.dθ̇_next_dy = 0.0\n",
    "        self.dθ̇_next_dθ = 0.0\n",
    "        self.dθ̇_next_dẏ = 0.0\n",
    "        self.dθ̇_next_dθ̇ = 0.0\n",
    "        self.dθ̇_next_du = 0.0\n",
    "        self.dy_next_dy = 0.0\n",
    "        self.dy_next_dθ = 0.0\n",
    "        self.dy_next_dẏ = 0.0\n",
    "        self.dy_next_dθ̇ = 0.0\n",
    "        self.dy_next_du = 0.0\n",
    "        self.dθ_next_dy = 0.0\n",
    "        self.dθ_next_dθ = 0.0\n",
    "        self.dθ_next_dẏ = 0.0\n",
    "        self.dθ_next_dθ̇ = 0.0\n",
    "        self.dθ_next_du = 0.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Validating Analytical Derivatives with Numerical Differentiation\n",
    "\n",
    "After defining the action model and deriving the analytical derivatives of the system’s dynamics and cost function, it is important to verify the correctness of these derivatives. One common method for validating derivatives is **numerical differentiation (NumDiff)**. This method approximates the derivative of a function by comparing the values of the function at small increments around a given point.\n",
    "\n",
    "The principle behind numerical differentiation is the approximation of the derivative of a function \\( f(x) \\) using the following formula:\n",
    "\n",
    "\\begin{equation}\n",
    "\\frac{df(x)}{dx} \\approx \\frac{f(x + \\epsilon) - f(x)}{\\epsilon}\n",
    "\\end{equation}\n",
    "\n",
    "where:\n",
    " - $\\epsilon$ is a small perturbation value (often called the step size), which controls the precision of the derivative approximation.\n",
    " - $f(x)$ is the function whose derivative is being approximated.\n",
    "\n",
    "In Crocoddyl, we perform numerical differentiation of an action model via `ActionModelNumDiff` class. This class is designed to define $\\epsilon$ values appropriately. This is done by consider the type of derivatives and floating-point number (`double`, `float`). Below, we show how to use for asserting our analytical derivatives."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create an action model, data, and its NumDiff's action\n",
    "model = ActionModelCartpole()\n",
    "data = model.createData()\n",
    "model_nd = crocoddyl.ActionModelNumDiff(model)\n",
    "data_nd = model_nd.createData()\n",
    "\n",
    "# Generate a random state and control\n",
    "x = model.state.rand()\n",
    "u = np.random.rand(1)\n",
    "\n",
    "# Compute analytical and NumDiff derivatives\n",
    "model.calc(data, x, u)\n",
    "model.calcDiff(data, x, u)\n",
    "model_nd.calc(data_nd, x, u)\n",
    "model_nd.calcDiff(data_nd, x, u)\n",
    "\n",
    "# Compare derivatives\n",
    "print(\"Fx difference:\", np.linalg.norm(data.Fx - data_nd.Fx))\n",
    "print(\"Fu difference:\", np.linalg.norm(data.Fu - data_nd.Fu))\n",
    "print(\"Lx difference:\", np.linalg.norm(data.Lx - data_nd.Lx))\n",
    "print(\"Lu difference:\", np.linalg.norm(data.Lu - data_nd.Lu))\n",
    "print(\"Lxx difference:\", np.linalg.norm(data.Lxx - data_nd.Lxx))\n",
    "print(\"Lxu difference:\", np.linalg.norm(data.Lxu - data_nd.Lxu))\n",
    "print(\"epsilon\", model_nd.disturbance)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simulating our Cart-Pole System\n",
    "\n",
    "After we have verified the analytical derivatives, we want to analyse/validate the cart-pole dynamics. To do so, we rollout the system dynamics without applying control commands. In Crocoddyl, this can be done easily using our `ShootingProblem` class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a shooting problem with 50 running nodes (i.e., 50 * 5e-2 = 2.5 sec)\n",
    "N = 50\n",
    "x0 = np.array([0.0, 0.0, 1.0, 0.5])\n",
    "problem = crocoddyl.ShootingProblem(x0, [model] * N, model)\n",
    "\n",
    "xs = problem.rollout([np.zeros(model.nu)] * N)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In `cartpole_utils`, we provide two functions for plotting and animating our cart-pole. These functions are `plotCartpole` and `animateCartpole`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import HTML\n",
    "from cartpole_utils import animateCartpole\n",
    "\n",
    "anim = animateCartpole(xs)\n",
    "\n",
    "# HTML(anim.to_jshtml())\n",
    "HTML(anim.to_html5_video())\n",
    "\n",
    "# If you encounter problems probably you need to install ffmpeg/libav-tools:\n",
    "# sudo apt-get install ffmpeg\n",
    "# or\n",
    "# sudo apt-get install libav-tools\n",
    "# or\n",
    "# conda install ffmpeg\n",
    "# Additionally, you need to install matplolib."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As expected, the cart is moving to the left, and the pole is falling down. This behavior occurs because the initial velocities of both the cart and the pole are non-zero. To stabilize the system and find an optimal trajectory for balancing the cart-pole, we formulate a nonlinear optimal control problem as follows:\n",
    "\n",
    "\\begin{aligned}\n",
    "& \\underset{u(t)}{\\text{minimize}} \\quad J = \\frac{1}{2} \\sum_{i=0}^{N} w_i \\, r_i^2(\\mathbf{x}, \\mathbf{u})  \\\\\n",
    "& \\text{subject to} \\quad \\mathbf{x}^+ = f(\\mathbf{x}, \\mathbf{u})\n",
    "\\end{aligned}\n",
    "\n",
    "Where:\n",
    " - $J$ is the cost function, which sums over the trajectory nodes  i  the weighted square of the residual  r_i(\\mathbf{x}, \\mathbf{u}) , capturing the deviation from the desired state and control input.\n",
    " - $w_i$ is the weight associated with each node in the trajectory, typically representing how much importance we assign to a given time step in the optimization.\n",
    " - $\\mathbf{x}^+$  represents the state at the next time step, while $\\mathbf{x}$ is the current state.\n",
    " - $f(\\mathbf{x}, \\mathbf{u})$ represents the system’s dynamics, describing how the state evolves from one time step to the next under the control input $\\mathbf{u}[i]$.\n",
    "\n",
    "This formulation defines the cost and dynamics of each node $i$ based on the cart-pole’s action model. All of these elements are encapsulated within Crocoddyl’s shooting problem framework, where the optimization problem is structured as a trajectory optimization problem.\n",
    "\n",
    "With this setup, we now have all the necessary components to solve the problem using our feasibility-driven Differential Dynamic Programming (DDP) solver. The DDP solver will iteratively adjust the control inputs $\\mathbf{u}[i]$ to minimize the cost function while ensuring the system adheres to its dynamics, ultimately finding an optimal trajectory that stabilizes the cart-pole system."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Creating the FDDP solver and setting the logger callback\n",
    "solver = crocoddyl.SolverFDDP(problem)\n",
    "solver.setCallbacks([crocoddyl.CallbackVerbose()])\n",
    "\n",
    "# Solving this problem and display the optimal trajectory of our cart-pole system\n",
    "solver.solve()\n",
    "anim = animateCartpole(solver.xs)\n",
    "\n",
    "# HTML(anim.to_jshtml())\n",
    "HTML(anim.to_html5_video())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can observe how the **Feasibily-Driven Differential Dynamic Programming (FDDP)** algorithm iteratively finds the optimal trajectory and control sequence. Additionally, the logger provides valuable insights at each iteration, displaying key quantities such as the cost, gradient, infeasibilities, and other relevant metrics. These quantities help us understand the numerical behavior of the optimization process, including convergence patterns and any potential issues during the search for the optimal solution.\n",
    "\n",
    "## What's Next?\n",
    "\n",
    "Now that we’ve explored the optimal control problem and seen how FDDP can be applied to find optimal trajectories and control sequences, the next step is to understand how action models can be code-generated for various systems. In this section, we will walk through:\n",
    " 1. **Action Model Code Generation**: Learn how to automatically generate the action model for cart-pole and other robotic systems using the `ActionModelCodeGen` class.\n",
    " 2. **Derivatives and Sensitivity Analysis**: Understand how to validate the action models by comparing the code-generated derivatives with numerical derivatives generated using tools like NumDiff.\n",
    " 3. **Inspecting Code-Generated Performance**: Familiarize yourself with Crocoddyl’s profiling tools to analyze the performance of the generated action models and evaluate computation efficiency.\n",
    " 4. **Scaling to Different Systems**: Once you’re familiar with the cart-pole system, we will show how to extend the code generation process to more complex robotic systems, enabling efficient computation and optimization.\n",
    "\n",
    "By the end of the `05_codegenerating_a_cartpole.ipynb`, you’ll be equipped to generate and validate action models for your own robotics applications and simulations, improving the efficiency and accuracy of your control systems. 🚀"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "crocoddyl",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
