tutorials/brb_brittle_star_cpg_tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "bad103a4db8ef575",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "# <h1><center>Central Pattern Generator Tutorial</center></h1>\n",
    "\n",
    "This notebook provides an introductory tutorial to the usage of Central Pattern Generators (CPG) for robotic locomotion. Specifically, we cover the bio-inspired origin of CPGs and how they can be implemented using differential equations. Furthermore, we show how these differential equations can be solved numerically using Jax.\n",
    "\n",
    "As a case study, we will define a CPG system and use it to imitate a brittle star's rowing gait. Therefore, we will use the [brittle Star simulation environment](https://github.com/Co-Evolve/brt/tree/main/biorobot/brittle_star), which is part of the [**the Bio-inspired Robotics Testbed (BRT)**](https://github.com/Co-Evolve/brt)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "848972e4d1afb3da",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Central Pattern Generators\n",
    "Central Pattern Generators (CPGs) are neural circuits that produce rhythmic patterned outputs without requiring rhythmic input. The term 'central' indicates that sensory feedback (from the peripheral nervous system) is not needed for generating the rhythms. This intrinsic capability to produce rhythmic patterns is essential for many basic life functions such as chewing, breathing, and digesting. They are also fundamental building blocks for the locomotor neural circuits both in invertebrate and vertebrate animals. CPGs present several interesting properties such as distributed control, the ability to deal with redundancies, fast control loops, and allowing modulation of the rhythmic patterns with simple control signals. These properties make CPGs interesting building blocks for locomotion controllers in robotics. \n",
    "\n",
    "Two primary approaches are prevalent in the implementation of Central Pattern Generators (CPGs) in robotics: the connectionist approach and the differential equations approach. Both rely on the creation of a nonlinear dynamical system.\n",
    "\n",
    "The connectionist approach involves creating CPGs through networks of interconnected artificial neurons (resulting in a recurrent neural network). The connectionist model is particularly effective for simulating adaptive behaviors, as it can learn and modify its connections based on feedback, much like a biological neural network. In essence, the connectionist approach relies on the Universal Approximation theorem of neural networks to provide a way to optimize towards any kind of rhythmic output. This approach is often employed when the targeted CPG properties cannot be modelled mathematically.\n",
    " \n",
    "On the other hand, the differential equations approach uses mathematical models based on differential equations to simulate the rhythmic outputs of CPGs. This method relies on carefully crafted equations that describe the dynamics of the system, often focusing on the interactions between different elements of the CPG and their temporal evolution. This approach is particularly valued for its precision and predictability, allowing for a detailed understanding of the underlying mechanics of rhythmic patterns. It is widely used in more traditional engineering and robotics contexts, where precise control and predictability are essential. In essence, we can enforce specific properties of rhythmic output we desire though crafting the differential equations.\n",
    "\n",
    "In this tutorial, we will focus on the latter approach."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed6139aa8b3a3c4f",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Differential equations\n",
    "Before diving into an actual CPG implementation, we will take a brief side step and recap what differential equations are and what they are used for.\n",
    "\n",
    "Differential equations are mathematical equations that describe the relationship between a function and its derivatives. They are fundamental in expressing how a quantity changes over time or space and are pivotal in modeling the behavior of diverse physical, biological, chemical, and economic systems. At their core, differential equations involve functions and their rates of change. For example, in physics, they are used to describe phenomena like the motion of a pendulum or the flow of heat in a solid. In biology, they can model population dynamics or the spread of diseases. In engineering, differential equations are crucial for designing control systems or understanding the behavior of electrical circuits. There are two main types of differential equations: ordinary differential equations (ODEs) and partial differential equations (PDEs). ODEs involve functions of a single variable and their derivatives, such as time. PDEs, on the other hand, deal with functions of multiple variables and their partial derivatives, like space and time. Solving a differential equation typically means finding the function or set of functions that satisfy the equation. The solutions can provide vital insights into the behavior of the system being modeled, predicting how it evolves over time or reacts to different conditions. Due to the complexity of many differential equations, solutions are often sought through numerical methods, especially for equations that cannot be solved analytically."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d8e288ae134f46b",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### Case study\n",
    "Given our use of physics simulators, we will use a related example. Physics simulators rely on differential equations to accurately represent the dynamics of real-world phenomena. Differential equations enable the simulation of a wide array of physical processes, such as the motion of objects under the influence of forces, fluid dynamics, and even changes in temperature and energy. The motion of objects is often governed by Newton's laws of motion, which are expressed through differential equations. For example, the change in velocity (i.e. the acceleration) of an object is calculated based on the forces applied to it, encapsulating the principles of mass, force, and acceleration. Similarly, in simulating fluid dynamics, the Navier-Stokes equations – a set of nonlinear partial differential equations – are used to describe the motion of fluid substances. These simulators numerically solve these differential equations at discrete time intervals. Ultimately, physics simulators thus solve differential equations to predict the state of the system at a future point in time, given the current state (i.e. the initial/current conditions of the system).\n",
    "\n",
    "In this case study, we will solve the differential equations that describe the motion of a point in one-dimensional space. As stated above, the motion of objects is often described by Newton's laws of motion: $F = ma$ (the second law). This law links the force $F$ we apply to an object to its mass $m$ and acceleration $a$. As you you know, acceleration $a$ is the derivative of velocity $v$ with respect to time $t$, which again is the derivative of position $p$ with respect to time. This law can thus be rewritten as $F = ma = m v' = m p''$, where $'$ denotes the derivative.\n",
    "\n",
    "To simplify things, let's say that in our one-dimensional space, we only have a single force that acts upon the object, which we completely control. We thus know $F(t)$. We will also give our object a mass of 1 such that we can remove it from the above equations: $F = a = v' = p''$. So given a function that describes our force over time $F(t)$, we immediately know what the acceleration $a$ of our point will be. But what if we want to know the velocity over time $v(t)$? Let's work this out with a force function $F(t) = t^2$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "initial_id",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "def force(\n",
    "        time: float\n",
    "        ) -> float:\n",
    "    return time ** 2\n",
    "\n",
    "\n",
    "def acceleration(\n",
    "        time: float,\n",
    "        velocity: float\n",
    "        ) -> float:\n",
    "    # F = ma\n",
    "    mass = 1\n",
    "    return mass * force(time=time)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae858c8d55dd033f",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Note: the reason why the acceleration function also takes the velocity as an argument, is because the general form of an ODE is $y'(t) = f(t, y(t))$, i.e. the derivative of a function is a function of the independent variable and the function itself. In our case, this means that $a(t) = v'(t) = f'(t, v(t)$. So while we do not explicitly use the velocity in the calculation of our acceleration, let's keep it as an argument anyway to stay general.\n",
    "\n",
    "Anyway, we know that $a = v'$, i.e. the derivative of our velocity with respect to time is equal to the acceleration. This means that the acceleration function describes how our velocity changes over time. If we were given the current velocity as an initial condition, i.e. $v(t_n)$, we can thus calculate $v(t_{n+1})$ by just taking a small steps in the direction of the acceleration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c68766440880083b",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def get_next_velocity(\n",
    "        current_time: float,\n",
    "        current_velocity: float,\n",
    "        delta_time: float\n",
    "        ) -> float:\n",
    "    current_acceleration = acceleration(time=current_time, velocity=current_velocity)\n",
    "    new_velocity = current_velocity + delta_time * current_acceleration\n",
    "    return new_velocity"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f79da3a5769a3e57",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Let's use this to simulate a trajectory of `simulation_time` seconds with our point, starting from $v(t_0) = 0$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2fd45bf2082e72eb",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "simulation_time = 10\n",
    "delta_time = 0.01\n",
    "\n",
    "times = [0]\n",
    "velocities = [0]\n",
    "\n",
    "while times[-1] < simulation_time:\n",
    "    next_velocity = get_next_velocity(current_time=times[-1], current_velocity=velocities[-1], delta_time=delta_time)\n",
    "    velocities.append(next_velocity)\n",
    "    times.append(times[-1] + delta_time)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f6a95e34284c962",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Time to write a visualization function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57b47cdc806a3a1a",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "\n",
    "\n",
    "def visualize_trajectory(\n",
    "        times: np.ndarray,\n",
    "        quantity: np.ndarray,\n",
    "        quantity_label: str,\n",
    "        color: str = \"blue\",\n",
    "        flush: bool = True\n",
    "        ) -> None:\n",
    "    plt.plot(times, quantity, color=color, label=quantity_label)\n",
    "    if flush:\n",
    "        plt.xlabel(\"time\")\n",
    "        plt.legend()\n",
    "        plt.show()\n",
    "        plt.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f361cdc1d69cab4",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Now we can visualize our velocity over time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a6bba9c8d68c46e",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "visualize_trajectory(times=times, quantity=velocities, quantity_label=\"velocity\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62ba867e1bde61f8",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "This is quite as expected. Our force function is $F(t) = t^2$. As discussed above, our point's acceleration $a$ follows this same function (given that we used a mass $m$ of 1). So we actually solved the ODE $a = v'$. This methodology, i.e. starting from a known initial condition and iteratively taking a small step in the direction of the derivative function (which actually represents the slope of the original function at that point) is known as the [Euler method](https://en.wikipedia.org/wiki/Euler_method).\n",
    "\n",
    "### Euler method\n",
    "Before moving on, let's first implement our Euler 'solver' in a more general manner."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "efa3e9a71944815e",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from typing import Callable\n",
    "\n",
    "\n",
    "def euler_solver(\n",
    "        current_time: float,\n",
    "        y: float,\n",
    "        derivative_fn: Callable[[float, float], float],\n",
    "        delta_time: float\n",
    "        ) -> float:\n",
    "    slope = derivative_fn(current_time, y)\n",
    "    next_y = y + delta_time * slope\n",
    "    return next_y"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20aade96d3d1532f",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "It's important to remember that we're actually approximating the velocity function numerically, instead of getting an exact solution. In this specific ODE where we have $v'(t) = a(t) = t^2$ and $v(0) = 0$, we can easily see that $v(t) = \\frac{1}{3}t^3$ is the exact solution (plug this into the ODE to check). See what happens if we try some different `delta_time` values (i.e. different step sizes) and compare it to this exact solution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81982ca6d260c607",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "simulation_time = 10\n",
    "delta_times = [0.1, 0.5, 1.0]\n",
    "colors = [\"red\", \"green\", \"blue\"]\n",
    "\n",
    "for delta_time, color in zip(delta_times, colors):\n",
    "    times = [0]\n",
    "    velocities = [0]\n",
    "\n",
    "    while times[-1] < simulation_time:\n",
    "        next_velocity = euler_solver(\n",
    "                current_time=times[-1], y=velocities[-1], derivative_fn=acceleration, delta_time=delta_time\n",
    "                )\n",
    "        velocities.append(next_velocity)\n",
    "        times.append(times[-1] + delta_time)\n",
    "    visualize_trajectory(\n",
    "            times=times, quantity=velocities, quantity_label=f\"delta_time {delta_time}\", color=color, flush=False\n",
    "            )\n",
    "\n",
    "exact_solution = 1 / 3 * np.power(np.array(times), 3)\n",
    "visualize_trajectory(times=times, quantity=exact_solution, quantity_label=\"exact velocity\", color=\"black\", flush=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aafab18f3b1823e7",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "We can see that the smaller our step size, the better the approximation! This is because we linearly approximate the change of our function $v(t)$ over time (by taking a step in the direction of our slope at each point), while the true change of our function (in this case $a(t) = t^2$) is non-linear! This approximation only remains sufficiently close to the true solution if we take small enough steps, at the cost of more computation.\n",
    "\n",
    "Sidenote: When you solve an ODE, you are essentially integrating the differential equation. The ODE describes the change over time, integrating this change over time (i.e. taking a 'sum' of all change over time), gives us the current 'state' at each time. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3919379d5988aeed",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "This Euler method is a member of the [Runge-Kutta methods](https://en.wikipedia.org/wiki/Runge%E2%80%93Kutta_methods). The Euler method is the simplest form of the Runge-Kutta methods and can be considered a first-order Runge-Kutta method. As shown above, it approximates the solution of an ODE by taking a series of steps along the curve, using the slope of the curve (derived from the ODE) at the beginning of each step to predict the value of the function at the next step. The process is akin to taking small, linear steps along the curve, each step based only on the immediate slope. In contrast, higher-order Runge-Kutta methods, such as the popular fourth-order Runge-Kutta method (RK4), take a more sophisticated approach. They use multiple evaluations of the slope (derivative) at various points within each step to make a more accurate prediction for the next value. These methods effectively use additional information about the curvature and behavior of the function within each step to refine their predictions. So at the cost of more function evaluations, it provides better slope approximations.\n",
    "\n",
    "### Runge-Kutta methods\n",
    "Let's implement an RK4 solver (4 refers to 4 slope approximations via 4 derivative function evaluations)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b67a7bdea0bfdd1",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def rk4_solver(\n",
    "        current_time: float,\n",
    "        y: float,\n",
    "        derivative_fn: Callable[[float, float], float],\n",
    "        delta_time: float\n",
    "        ) -> float:\n",
    "    # This is the original euler\n",
    "    slope1 = derivative_fn(current_time, y)\n",
    "    # These are additional slope calculations that improve our approximation of the true slope  \n",
    "    slope2 = derivative_fn(current_time + delta_time / 2, y + slope1 * delta_time / 2)\n",
    "    slope3 = derivative_fn(current_time + delta_time / 2, y + slope2 * delta_time / 2)\n",
    "    slope4 = derivative_fn(current_time + delta_time, y + slope3 * delta_time)\n",
    "    average_slope = (slope1 + 2 * slope2 + 2 * slope3 + slope4) / 6\n",
    "    next_y = y + average_slope * delta_time\n",
    "    return next_y"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27ffb029bdf7b401",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Let's compare the euler and RK4 approximations!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "804762e3f2e045e9",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "simulation_time = 10\n",
    "delta_time = 0.5\n",
    "solvers = [euler_solver, rk4_solver]\n",
    "solver_names = [\"euler\", \"rk4\"]\n",
    "colors = [\"red\", \"blue\"]\n",
    "\n",
    "for solver, solver_name, color in zip(solvers, solver_names, colors):\n",
    "    times = [0]\n",
    "    velocities = [0]\n",
    "\n",
    "    while times[-1] < simulation_time:\n",
    "        next_velocity = solver(\n",
    "                current_time=times[-1], y=velocities[-1], derivative_fn=acceleration, delta_time=delta_time\n",
    "                )\n",
    "        velocities.append(next_velocity)\n",
    "        times.append(times[-1] + delta_time)\n",
    "    visualize_trajectory(\n",
    "            times=times, quantity=velocities, quantity_label=f\"solver: {solver_name}\", color=color, flush=False\n",
    "            )\n",
    "\n",
    "exact_solution = 1 / 3 * np.power(np.array(times), 3)\n",
    "visualize_trajectory(times=times, quantity=exact_solution, quantity_label=\"exact velocity\", color=\"black\", flush=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7899344cb699ba6c",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "As you can see, our RK4 solver solves the ODE better than our orignal Euler (which could also be written as RK1). We can even say that the RK4 solver provides an exact solution to the ODE because if the exact solution to the differential equation is a polynomial of order $n$, it will be solved exactly by an $n$-th order Runge-Kutta method. The slope approximations that appear in the RK$n$ methods are based on a Taylor series approximation of the true slope. Diving into that is a bit out of scope for this tutorial, but feel free to head to dive a little bit deeper yourself."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8709f848c7c627c5",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### Second-order differential equations \n",
    "The velocity and acceleration example from above is an example of a first-order ODE (i.e. the highest order of a derivative in the ODE is one). We used the euler and RK4 methods to solve for the velocity $v(t)$ given an acceleration function $a(t)$. If we now want to find our point's position $p(t)$ over time, we get the following second-order ODE: $a(t) = p''(t)$. The nice thing is that we can reuse or previous solution to do so! Since we already solved the $a(t) = v'(t)$, we know the (approximated) value of the velocity function $v(t)$. So we can just solve another first order ODE instead: $v(t) = p'(t)$!\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd63947170ef1563",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "times = [0]\n",
    "positions = [0]\n",
    "velocities = [0]\n",
    "delta_time = 0.01\n",
    "\n",
    "while times[-1] < 10:\n",
    "    next_velocity = rk4_solver(\n",
    "            current_time=times[-1], y=velocities[-1], derivative_fn=acceleration, delta_time=delta_time\n",
    "            )\n",
    "    next_position = rk4_solver(\n",
    "            current_time=times[-1],\n",
    "            y=positions[-1],\n",
    "            derivative_fn=lambda\n",
    "                t,\n",
    "                p: velocities[-1],\n",
    "            delta_time=delta_time\n",
    "            )\n",
    "    times.append(times[-1] + delta_time)\n",
    "    positions.append(next_position)\n",
    "    velocities.append(next_velocity)\n",
    "\n",
    "visualize_trajectory(times=times, quantity=positions, quantity_label=\"position\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b083a0626f119af5",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Knowing that $v(t) = \\frac{1}{3}t^3$ and that $v(0) = 0$, we can again see that the exact solution would be $p(t) = \\frac{1}{3 * 4}t^4$. Let's compare our previous result to the exact solution as a validation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20024ac90acb0fb6",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "visualize_trajectory(times=times, quantity=positions, quantity_label=\"approximated position\", flush=False)\n",
    "exact_solution = 1 / 12 * np.power(times, 4)\n",
    "visualize_trajectory(times=times, quantity=exact_solution, quantity_label=\"true position\", color=\"black\", flush=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7dd0dfc9486efe2f",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Solving a second-order ODE thus comes down to solving two first-order ODEs. We solved $p'' = a = t^2$ with initial conditions $p(0) = 0, p'(0) = v(0) = 0$. We did this by first solving $v'(t) = t^2$ with the initial condition $v(0) = 0$. Every step, this gives us $v(t_{n+1})$\n",
    "Then we solve $p(t)' = v(t)$ with the initial condition $p(0) = 0$. In every step, we approximate $v(t)$ as a constant, given that we calculated it in the previous step (or know it as an initial condition due to $p'(0) = v(0)$ being given).\n",
    "\n",
    "In a more general form, let's say we we have a second-order ODE: $y'' + y' + y = c$  with initial conditions $y(0) = a, y'(0) = b$\n",
    "If we let $y_1 = y$ and $y_2 = y'$, we can write the original second-order ODE as a system of two first-order ODEs:\n",
    "$\\begin{bmatrix} y_1' \\\\ y_2'\\end{bmatrix} = \\begin{bmatrix} y_2 \\\\ c - y_1 - y_2\\end{bmatrix}$ with initial conditions: $\\begin{bmatrix} y_1(0) \\\\ y_2(0)\\end{bmatrix} = \\begin{bmatrix} a \\\\ b\\end{bmatrix}$. In this manner, we can use our first-order ODE solvers to solve second-order ODEs (and even higher order ODEs) by decomposing them."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6bbd110e2dda7b07",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Note: although our case study focused on solving ordinary differential equations, the same methodology can be employed for partial differential equations."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df9d409521eca365",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Implementing a Central Pattern Generator"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f3708526d00b9c7",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Now that we refreshed our intuition with respect to using and solving differential equations, we can finally move on to the actual implementation of a CPG. We will implement the CPG model used in a paper by [Sproewitz et al](https://journals.sagepub.com/doi/10.1177/0278364907088401). This CPG model is implemented as a system of $N$ coupled amplitude- and offset-controlled phase oscillators. An oscillator $i$ is defined by the following differential equations (the $\\dot{}$ refers to the derivative operator, which we previously used $`$ for): \n",
    "\n",
    "(1) $\\dot{\\phi_i} = \\omega_i + \\sum_j{w_{ij}r_j}\\sin{(\\phi_j - \\phi_i - \\rho_{ij})}$<br>\n",
    "(2) $\\ddot{r_i} = a_r(\\frac{a_r}{4}(R_i - r_i) - \\dot{r_i})$<br>\n",
    "(3) $\\ddot{x_i} = a_x(\\frac{a_x}{4}(X_i - x_i) - \\dot{x_i})$<br>\n",
    "(4) $\\theta_i = x_i + r_i\\cos{(\\phi_i)}$<br>\n",
    "\n",
    "where $\\theta_i$ is the output of the oscillator (in radians) and $\\phi_i$ , $r_i$ and $x_i$ are state variables that respectively encode the phase, amplitude and the offset of the oscillations. The parameters $\\omega_i$ , $R_i$ and $X_i$ are control parameters for the desired frequency, amplitude and offset of the oscillations. The parameters $w_{ij}$ and $\\rho_{ij}$ are respectively the coupling weights and phase biases, which determine how oscillator $j$ influences oscillator $i$. The parameters $a_r$ and $a_x$ are constant positive gains.\n",
    "\n",
    "These equations were designed such that the output of the oscillator $i$ exhibits limit cycle behavior, i.e. produces a stable periodic output. Equation (1) determines the time evolution of the phases of the oscillators, i.e. the phase frequency (\"in which direction and how fast should the phase change?\"). The sum term represents the coupling with other oscillators.   \n",
    "\n",
    "Similar to the paper, we will use the same frequency parameter $\\omega_i = \\omega$ for all oscillators, and bi-directional couplings between oscillators such that $\\rho_{ij} = -\\rho_{ji}$. We also ensure that, in case of a closed loop of inter-oscillator couplings, all phase biases in the loop are consistent (i.e. sum to a multiple of $2\\pi$). With these parameters, the phases will converge to a regime in which they grow linearly with a common rate $\\omega$ and with a phase difference between oscillators determined by $\\rho_{ij}$. Equations (2) and (3) are critically damped second-order [linear differential equations](https://en.wikipedia.org/wiki/Linear_differential_equation) which respectively have $R_i$ and $X_i$ as stable fixed points. This means that from any initial conditions, the state variables $r_i$ and $x_i$ will asymptotically and monotonically converge to $R_i$ and $X_i$. This allows one to smoothly modulate the amplitude and offset of oscillations.\n",
    "\n",
    "With these settings, two oscillators $i$ and $j$ that are coupled with non-zero weights $w_{ij}$ asymptotically converge to limit cycles (i.e. a closed trajectory in the state space that corresponds to sustained oscillations without decay or growth) $\\Theta_i^\\infty$ and $\\Theta_j^\\infty$ with the following closed form (i.e. exact) solutions:\n",
    "\n",
    "$\\Theta_i^\\infty = X_i + R_i\\cos{(\\omega t + \\rho_{ij} + \\phi_0)}$<br>\n",
    "$\\Theta_j^\\infty = X_j + R_j\\cos{(\\omega t + \\rho_{ji} + \\phi_0)}$<br>\n",
    "where $\\phi_0$ depends on the initial conditions of the system. Note that, because of the common frequency $\\omega$ and phase biases $\\phi_{ij}$ in a loop, this behaviour does not depend on the coupling weights $w_{ij}$. The coupling weights only affect how quickly the system converges: the larger the weights, the faster the convergence. This means that the system stabilizes into oscillations that are phase-locked for all oscillators that are connected together. These oscillations can be modulated (i.e. controlled) by several control parameters, $\\omega$ for setting a common frequency, $\\rho_{ij}$ for setting the phase lags between two connected oscillators, $R_i$ for setting the individual amplitudes and $X_i$ for setting the individual offsets.\n",
    "\n",
    "This CPG model has several nice properties:\n",
    "1. The system exhibits limit cycle behavior, i.e. oscillations rapidly return to the steady-state oscillations after any transient perturbation of the state variables\n",
    "2. This limit cycle has a closed form solution. Most types of oscillators used to build CPGs do not have a closed form solution for their limit cycle. The limit cycle has a harmonic shape and has control parameters ($\\omega$, $R_i$ and $X_i$) that are explicit and are directly related to relevant features of the oscillations. This facilitates the design and analysis of locomotion controllers. \n",
    "3. A third interesting property is that these control parameters can be abruptly and/or continuously varied while inducing only smooth changes in the resulting oscillations (i.e. there are no discontinuities nor jerks). This property is important in avoiding damage in the motors and gearboxes of robots.\n",
    "\n",
    "Let's now implement this model for a system with 2 bi-directionally coupled oscillators in JAX. We will implement this using a functional CPG class and a datastructures for the CPG's state. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cbb3f19d16adfe16",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import functools\n",
    "from flax import struct\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import chex\n",
    "\n",
    "\n",
    "@struct.dataclass\n",
    "class CPGState:\n",
    "    time: float\n",
    "    phases: jnp.ndarray\n",
    "    dot_amplitudes: jnp.ndarray  # first order derivative of the amplitude\n",
    "    amplitudes: jnp.ndarray\n",
    "    dot_offsets: jnp.ndarray  # first order derivative of the offset \n",
    "    offsets: jnp.ndarray\n",
    "    outputs: jnp.ndarray\n",
    "\n",
    "    # We'll make these modulatory parameters part of the state as they will change as well\n",
    "    R: jnp.ndarray\n",
    "    X: jnp.ndarray\n",
    "    omegas: jnp.ndarray\n",
    "    rhos: jnp.ndarray\n",
    "\n",
    "\n",
    "class CPG:\n",
    "    def __init__(\n",
    "            self,\n",
    "            weights: jnp.ndarray,\n",
    "            amplitude_gain: float = 20,\n",
    "            offset_gain: float = 20,\n",
    "            dt: float = 0.01,\n",
    "            solver: str = \"euler\"\n",
    "            ) -> None:\n",
    "        self._weights = weights\n",
    "        self._amplitude_gain = amplitude_gain\n",
    "        self._offset_gain = offset_gain\n",
    "        self._dt = dt\n",
    "        assert solver in [\"euler\", \"rk4\"], f\"'solver' must be one of ['euler', 'rk4']\"\n",
    "\n",
    "        if solver == \"euler\":\n",
    "            self._solver = euler_solver\n",
    "        else:\n",
    "            self._solver = rk4_solver\n",
    "\n",
    "    @property\n",
    "    def num_oscillators(\n",
    "            self\n",
    "            ) -> int:\n",
    "        return self._weights.shape[0]\n",
    "\n",
    "    @staticmethod\n",
    "    def phase_de(\n",
    "            weights: jnp.ndarray,\n",
    "            amplitudes: jnp.ndarray,\n",
    "            phases: jnp.ndarray,\n",
    "            phase_biases: jnp.ndarray,\n",
    "            omegas: jnp.ndarray\n",
    "            ) -> jnp.ndarray:\n",
    "        @jax.vmap  # vectorizes this function for us over an additional batch dimension (in this case over all oscillators)\n",
    "        def sine_term(\n",
    "                phase_i: float,\n",
    "                phase_biases_i: float\n",
    "                ) -> jnp.ndarray:\n",
    "            return jnp.sin(phases - phase_i - phase_biases_i)\n",
    "\n",
    "        couplings = jnp.sum(weights * amplitudes * sine_term(phase_i=phases, phase_biases_i=phase_biases), axis=1)\n",
    "        return omegas + couplings\n",
    "\n",
    "    @staticmethod\n",
    "    def second_order_de(\n",
    "            gain: jnp.ndarray,\n",
    "            modulator: jnp.ndarray,\n",
    "            values: jnp.ndarray,\n",
    "            dot_values: jnp.ndarray\n",
    "            ) -> jnp.ndarray:\n",
    "\n",
    "        return gain * ((gain / 4) * (modulator - values) - dot_values)\n",
    "\n",
    "    @staticmethod\n",
    "    def first_order_de(\n",
    "            dot_values: jnp.ndarray\n",
    "            ) -> jnp.ndarray:\n",
    "        return dot_values\n",
    "\n",
    "    @staticmethod\n",
    "    def output(\n",
    "            offsets: jnp.ndarray,\n",
    "            amplitudes: jnp.ndarray,\n",
    "            phases: jnp.ndarray\n",
    "            ) -> jnp.ndarray:\n",
    "        return offsets + amplitudes * jnp.cos(phases)\n",
    "\n",
    "    def reset(\n",
    "            self,\n",
    "            rng: chex.PRNGKey\n",
    "            ) -> CPGState:\n",
    "        phase_rng, amplitude_rng, offsets_rng = jax.random.split(rng, 3)\n",
    "        # noinspection PyArgumentList\n",
    "        state = CPGState(\n",
    "                phases=jax.random.uniform(\n",
    "                        key=phase_rng, shape=(self.num_oscillators,), dtype=jnp.float32, minval=-0.01, maxval=0.01\n",
    "                        ),\n",
    "                amplitudes=jnp.zeros(self.num_oscillators),\n",
    "                offsets=jnp.zeros(self.num_oscillators),\n",
    "                dot_amplitudes=jnp.zeros(self.num_oscillators),\n",
    "                dot_offsets=jnp.zeros(self.num_oscillators),\n",
    "                outputs=jnp.zeros(self.num_oscillators),\n",
    "                time=0.0,\n",
    "                R=jnp.zeros(self.num_oscillators),\n",
    "                X=jnp.zeros(self.num_oscillators),\n",
    "                omegas=jnp.zeros(self.num_oscillators),\n",
    "                rhos=jnp.zeros_like(self._weights)\n",
    "                )\n",
    "        return state\n",
    "\n",
    "    @functools.partial(jax.jit, static_argnums=(0,))\n",
    "    def step(\n",
    "            self,\n",
    "            state: CPGState\n",
    "            ) -> CPGState:\n",
    "        # Update phase\n",
    "        new_phases = self._solver(\n",
    "                current_time=state.time,\n",
    "                y=state.phases,\n",
    "                derivative_fn=lambda\n",
    "                    t,\n",
    "                    y: self.phase_de(\n",
    "                        omegas=state.omegas,\n",
    "                        amplitudes=state.amplitudes,\n",
    "                        phases=y,\n",
    "                        phase_biases=state.rhos,\n",
    "                        weights=self._weights\n",
    "                        ),\n",
    "                delta_time=self._dt\n",
    "                )\n",
    "        new_dot_amplitudes = self._solver(\n",
    "                current_time=state.time,\n",
    "                y=state.dot_amplitudes,\n",
    "                derivative_fn=lambda\n",
    "                    t,\n",
    "                    y: self.second_order_de(\n",
    "                        gain=self._amplitude_gain, modulator=state.R, values=state.amplitudes, dot_values=y\n",
    "                        ),\n",
    "                delta_time=self._dt\n",
    "                )\n",
    "        new_amplitudes = self._solver(\n",
    "                current_time=state.time,\n",
    "                y=state.amplitudes,\n",
    "                derivative_fn=lambda\n",
    "                    t,\n",
    "                    y: self.first_order_de(dot_values=state.dot_amplitudes),\n",
    "                delta_time=self._dt\n",
    "                )\n",
    "        new_dot_offsets = self._solver(\n",
    "                current_time=state.time,\n",
    "                y=state.dot_offsets,\n",
    "                derivative_fn=lambda\n",
    "                    t,\n",
    "                    y: self.second_order_de(\n",
    "                        gain=self._offset_gain, modulator=state.X, values=state.offsets, dot_values=y\n",
    "                        ),\n",
    "                delta_time=self._dt\n",
    "                )\n",
    "        new_offsets = self._solver(\n",
    "                current_time=0,\n",
    "                y=state.offsets,\n",
    "                derivative_fn=lambda\n",
    "                    t,\n",
    "                    y: self.first_order_de(dot_values=state.dot_offsets),\n",
    "                delta_time=self._dt\n",
    "                )\n",
    "\n",
    "        new_outputs = self.output(offsets=new_offsets, amplitudes=new_amplitudes, phases=new_phases)\n",
    "        # noinspection PyUnresolvedReferences\n",
    "        return state.replace(\n",
    "                phases=new_phases,\n",
    "                dot_amplitudes=new_dot_amplitudes,\n",
    "                amplitudes=new_amplitudes,\n",
    "                dot_offsets=new_dot_offsets,\n",
    "                offsets=new_offsets,\n",
    "                outputs=new_outputs,\n",
    "                time=state.time + self._dt\n",
    "                )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5eee6e2b8966d58",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Time to test! To validate the CPG implementation, it's always a good idea to visualize its output. We'll do that here, and we'll also visualize the phases of the two oscillators. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ceb1f78f52f64002",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "rng = jax.random.PRNGKey(seed=0)\n",
    "\n",
    "adjacency_matrix = jnp.ones((2, 2)) - jnp.identity(2)\n",
    "weights = 5 * adjacency_matrix\n",
    "cpg = CPG(weights=weights, amplitude_gain=20, offset_gain=20, dt=0.01, solver=\"euler\")\n",
    "\n",
    "cpg_state = cpg.reset(rng=rng)\n",
    "\n",
    "# Set modulation parameters\n",
    "cpg_state = cpg_state.replace(\n",
    "        R=jnp.array([1.0, 0.5]),\n",
    "        X=jnp.array([1.0, -1.0]),\n",
    "        omegas=(jnp.pi / 2) * jnp.ones(cpg.num_oscillators),\n",
    "        rhos=adjacency_matrix * jnp.array([[jnp.pi], [-jnp.pi]])\n",
    "        )\n",
    "\n",
    "SIMLATION_TIME = 20\n",
    "\n",
    "# We'll save the states to visualize them at the end\n",
    "states = []\n",
    "while cpg_state.time < SIMLATION_TIME:\n",
    "    cpg_state = cpg.step(\n",
    "            cpg_state\n",
    "            )\n",
    "    states.append(cpg_state)\n",
    "\n",
    "times = [state.time for state in states]\n",
    "o1 = [state.outputs[0] for state in states]\n",
    "o2 = [state.outputs[1] for state in states]\n",
    "\n",
    "plt.plot(times, o1, color=\"blue\", label=\"oscillator 1: output\")\n",
    "plt.plot(times, o2, color=\"red\", label=\"oscillator 2: output\")\n",
    "plt.legend()\n",
    "plt.xlabel(\"time\")\n",
    "plt.ylabel(\"output\")\n",
    "plt.show()\n",
    "plt.close()\n",
    "\n",
    "times = [state.time for state in states]\n",
    "o1 = [state.phases[0] for state in states]\n",
    "o2 = [state.phases[1] for state in states]\n",
    "\n",
    "plt.plot(times, o1, color=\"blue\", label=\"oscillator 1: phase\")\n",
    "plt.plot(times, o2, color=\"red\", label=\"oscillator 2: phase\")\n",
    "plt.plot(times, np.array(o2) - np.array(o1), label=\"phase difference\", color=\"black\")\n",
    "plt.legend()\n",
    "plt.xlabel(\"time\")\n",
    "plt.ylabel(\"output\")\n",
    "plt.show()\n",
    "plt.close()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c7b44d613092819",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Cool! Looking at the CPG outputs, we can indeed see that our oscillators (1) converge to a limit cycle (steady-state oscillations), even though we have random initial conditions, (2) the two oscillators indeed converge to the given target amplitudes $R_0=1$ and $R_1=0.5$, (3) the two oscillators indeed converge to the given target offsets $X_0=1$ and $X_1=-1$, and (4)  the two oscillators indeed exhibit our requested phase lag $\\pi$, making them run in anti-phase. The latter is also demonstrated by the phases plot. We can see that the phase difference between the two oscillators converges to our phase bias setting: $\\pi$, i.e. $\\rho_{12} = \\pi$ and $\\rho_{21} = -\\pi$. Due to the definition of the oscillator couplings in Equation $(1)$ above, this causes the oscillator one to decrease its phase frequency, and oscillator two to increase its phase frequency, until the desired bias (i.e. difference between them) of $\\pi$ is met. \n",
    "\n",
    "Let's now see how our CPG reacts to intermediate modulations. What if we were to swap our amplitude and offset modulation parameters for both oscillators at $t=10$?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b2aeba14644f44b",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "rng = jax.random.PRNGKey(seed=0)\n",
    "\n",
    "adjacency_matrix = jnp.ones((2, 2)) - jnp.identity(2)\n",
    "weights = 5 * adjacency_matrix\n",
    "cpg = CPG(weights=weights, amplitude_gain=20, offset_gain=20, dt=0.01, solver=\"euler\")\n",
    "\n",
    "cpg_state = cpg.reset(rng=rng)\n",
    "\n",
    "# Set modulation parameters\n",
    "cpg_state = cpg_state.replace(\n",
    "        R=jnp.array([1.0, 0.5]),\n",
    "        X=jnp.array([1.0, -1.0]),\n",
    "        omegas=(jnp.pi / 2) * jnp.ones(cpg.num_oscillators),\n",
    "        rhos=adjacency_matrix * jnp.array([[jnp.pi], [-jnp.pi]])\n",
    "        )\n",
    "\n",
    "SIMLATION_TIME = 20\n",
    "\n",
    "# We'll save the states to visualize them at the end\n",
    "states = []\n",
    "remodulated = False\n",
    "while cpg_state.time < SIMLATION_TIME:\n",
    "    cpg_state = cpg.step(\n",
    "            cpg_state\n",
    "            )\n",
    "    states.append(cpg_state)\n",
    "\n",
    "    if cpg_state.time > 10 and not remodulated:\n",
    "        remodulated = True\n",
    "        cpg_state = cpg_state.replace(\n",
    "                R=jnp.array([0.5, 1.0]), X=jnp.array([-1.0, 1.0]), )\n",
    "\n",
    "times = [state.time for state in states]\n",
    "o1 = [state.outputs[0] for state in states]\n",
    "o2 = [state.outputs[1] for state in states]\n",
    "\n",
    "plt.plot(times, o1, color=\"blue\", label=\"oscillator 1: output\")\n",
    "plt.plot(times, o2, color=\"red\", label=\"oscillator 2: output\")\n",
    "plt.legend()\n",
    "plt.xlabel(\"time\")\n",
    "plt.ylabel(\"output\")\n",
    "plt.show()\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3e1bbc2c6d7349d",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Cheers! Take some time to play around with these parameters now. A good way to get a better intuition into the usage of CPGs is to try to predict the influence of some parameter changes before running them. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ff67868de35169c",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Using the CPG to imitate the brittle star's rowing behaviour\n",
    "\n",
    "Using a CPG for robot control comes down to defining a system of interconnected oscillators of which we pass the outputs to the actuators of the robot. Let's try to do this with our brittle star robot to imitate a real brittle star's rowing gait."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b0ea84fab771ff",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import subprocess\n",
    "import logging\n",
    "\n",
    "try:\n",
    "    if subprocess.run('nvidia-smi').returncode:\n",
    "        raise RuntimeError(\n",
    "                'Cannot communicate with GPU. '\n",
    "                'Make sure you are using a GPU Colab runtime. '\n",
    "                'Go to the Runtime menu and select Choose runtime type.'\n",
    "                )\n",
    "\n",
    "    # Add an ICD config so that glvnd can pick up the Nvidia EGL driver.\n",
    "    # This is usually installed as part of an Nvidia driver package, but the Colab\n",
    "    # kernel doesn't install its driver via APT, and as a result the ICD is missing.\n",
    "    # (https://github.com/NVIDIA/libglvnd/blob/master/src/EGL/icd_enumeration.md)\n",
    "    NVIDIA_ICD_CONFIG_PATH = '/usr/share/glvnd/egl_vendor.d/10_nvidia.json'\n",
    "    if not os.path.exists(NVIDIA_ICD_CONFIG_PATH):\n",
    "        with open(NVIDIA_ICD_CONFIG_PATH, 'w') as f:\n",
    "            f.write(\n",
    "                    \"\"\"{\n",
    "                            \"file_format_version\" : \"1.0.0\",\n",
    "                            \"ICD\" : {\n",
    "                                \"library_path\" : \"libEGL_nvidia.so.0\"\n",
    "                            }\n",
    "                        }\n",
    "                        \"\"\"\n",
    "                    )\n",
    "\n",
    "    # Configure MuJoCo to use the EGL rendering backend (requires GPU)\n",
    "    print('Setting environment variable to use GPU rendering:')\n",
    "    %env MUJOCO_GL=egl\n",
    "\n",
    "    # xla_flags = os.environ.get('XLA_FLAGS', '')\n",
    "    # xla_flags += ' --xla_gpu_triton_gemm_any=True'\n",
    "    # os.environ['XLA_FLAGS'] = xla_flags\n",
    "\n",
    "    # Check if jax finds the GPU\n",
    "    import jax\n",
    "\n",
    "    print(jax.devices('gpu'))\n",
    "except Exception:\n",
    "    logging.warning(\"Failed to initialize GPU. Everything will run on the cpu.\")\n",
    "\n",
    "try:\n",
    "    print('Checking that the mujoco installation succeeded:')\n",
    "    import mujoco\n",
    "\n",
    "    mujoco.MjModel.from_xml_string('<mujoco/>')\n",
    "except Exception as e:\n",
    "    raise e from RuntimeError(\n",
    "            'Something went wrong during installation. Check the shell output above '\n",
    "            'for more information.\\n'\n",
    "            'If using a hosted Colab runtime, make sure you enable GPU acceleration '\n",
    "            'by going to the Runtime menu and selecting \"Choose runtime type\".'\n",
    "            )\n",
    "\n",
    "print('MuJoCo installation successful.')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e93fa5ea8096c719",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### Environment setup\n",
    "We'll start by creating our environment. We'll use a brittle star robot with 5 arms and 5 segments per arm."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12c122bd5879da04",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from moojoco.environment.base import MuJoCoEnvironmentConfiguration\n",
    "from biorobot.brittle_star.environment.undirected_locomotion.dual import BrittleStarUndirectedLocomotionEnvironment\n",
    "from biorobot.brittle_star.environment.undirected_locomotion.shared import \\\n",
    "    BrittleStarUndirectedLocomotionEnvironmentConfiguration\n",
    "from typing import List\n",
    "import mediapy as media\n",
    "from biorobot.brittle_star.mjcf.morphology.morphology import MJCFBrittleStarMorphology\n",
    "from biorobot.brittle_star.mjcf.morphology.specification.default import default_brittle_star_morphology_specification\n",
    "from biorobot.brittle_star.mjcf.arena.aquarium import AquariumArenaConfiguration, MJCFAquariumArena\n",
    "\n",
    "morphology_specification = default_brittle_star_morphology_specification(\n",
    "        num_arms=5, num_segments_per_arm=5, use_p_control=True, use_torque_control=False\n",
    "        )\n",
    "arena_configuration = AquariumArenaConfiguration(\n",
    "        size=(10, 5), sand_ground_color=False, attach_target=False, wall_height=1.5, wall_thickness=0.1\n",
    "        )\n",
    "environment_configuration = BrittleStarUndirectedLocomotionEnvironmentConfiguration(\n",
    "        joint_randomization_noise_scale=0.0,\n",
    "        render_mode=\"rgb_array\",\n",
    "        simulation_time=20,\n",
    "        num_physics_steps_per_control_step=10,\n",
    "        time_scale=2,\n",
    "        camera_ids=[0, 1],\n",
    "        render_size=(480, 640)\n",
    "        )\n",
    "\n",
    "\n",
    "def create_environment() -> BrittleStarUndirectedLocomotionEnvironment:\n",
    "    morphology = MJCFBrittleStarMorphology(\n",
    "            specification=morphology_specification\n",
    "            )\n",
    "    arena = MJCFAquariumArena(\n",
    "            configuration=arena_configuration\n",
    "            )\n",
    "    env = BrittleStarUndirectedLocomotionEnvironment.from_morphology_and_arena(\n",
    "            morphology=morphology, arena=arena, configuration=environment_configuration, backend=\"MJX\"\n",
    "            )\n",
    "    return env\n",
    "\n",
    "\n",
    "def post_render(\n",
    "        render_output: List[np.ndarray],\n",
    "        environment_configuration: MuJoCoEnvironmentConfiguration\n",
    "        ) -> np.ndarray:\n",
    "    if render_output is None:\n",
    "        # Temporary workaround until https://github.com/google-deepmind/mujoco/issues/1379 is fixed\n",
    "        return None\n",
    "\n",
    "    num_cameras = len(environment_configuration.camera_ids)\n",
    "    num_envs = len(render_output) // num_cameras\n",
    "\n",
    "    if num_cameras > 1:\n",
    "        # Horizontally stack frames of the same environment\n",
    "        frames_per_env = np.array_split(render_output, num_envs)\n",
    "        render_output = [np.concatenate(env_frames, axis=1) for env_frames in frames_per_env]\n",
    "\n",
    "    # Vertically stack frames of different environments\n",
    "    render_output = np.concatenate(render_output, axis=0)\n",
    "\n",
    "    return render_output[:, :, ::-1]  # RGB to BGR\n",
    "\n",
    "\n",
    "def show_video(\n",
    "        images: List[np.ndarray | None]\n",
    "        ) -> str | None:\n",
    "    # Temporary workaround until https://github.com/google-deepmind/mujoco/issues/1379 is fixed\n",
    "    filtered_images = [image for image in images if image is not None]\n",
    "    num_nones = len(images) - len(filtered_images)\n",
    "    if num_nones > 0:\n",
    "        logging.warning(\n",
    "                f\"env.render produced {num_nones} None's. Resulting video might be a bit choppy (consquence of https://github.com/google-deepmind/mujoco/issues/1379).\"\n",
    "                )\n",
    "    return media.show_video(images=filtered_images)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ee45af954ab562b",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Let's create the environment.\n",
    "Since we're using an MJX environment, we'll also jit the `step` and `reset` functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b65aca465ec06cb8",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "env = create_environment()\n",
    "jit_step = jax.jit(env.step)\n",
    "jit_reset = jax.jit(env.reset)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85491fe2d3c76299",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Let's see what this environment looks like again, which observations it returns and which actions it expects.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9bfac0e74c2cc30",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "print(\"Observation space:\")\n",
    "print(f\"\\t{env.observation_space}\")\n",
    "print(\"Action space:\")\n",
    "print(f\"\\t{env.action_space}\")\n",
    "print(\"First 8 Actuators:\")\n",
    "print(f\"\\t{env.actuators[:8]}\")\n",
    "rng = jax.random.PRNGKey(0)\n",
    "state = jit_reset(rng=rng)\n",
    "media.show_image(post_render(env.render(state=state), environment_configuration=environment_configuration))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7071ac972d042cdd",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "This brittle star robot has $50$ joints in total (2 per segment: one in-plane (IP) joint and one out-of-plane (OOP) joint), which all have the same range of motion (as demonstrated by `env.action_space`). Each joint is given an actuator that can control the joints position. For the in-plane joints, a positive position means that the segment will point leftwards (when looking at the segment from the central disc), and rightwards for a negative position. For the out-of-plane joints, a positive position means that the segment will point upwards and downwards for a negative position. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f65848c1c535f8a5",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### Generating circular motion\n",
    "To produce simple rowing behaviour, we'll need to make the arms move in a circular motion. The simplest way to do so, would be to actuate the IP joint according to a cosine function, and the OOP joint according to a sine function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e610ff55b2b38f2c",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "times = np.linspace(0, 2 * np.pi, 100)\n",
    "in_plane_positions = np.cos(times)\n",
    "out_of_plane_positions = np.sin(times)\n",
    "\n",
    "plt.scatter(in_plane_positions, out_of_plane_positions, c=times, cmap=\"Blues\")\n",
    "plt.xlabel(\"in-plane position\")\n",
    "plt.ylabel(\"out-of-plane position\")\n",
    "plt.show()\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8dfff37c042abf0",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "This plot shows traversed IP and OOP positions (the time dimension is depicted by increasing brightness). \n",
    "Our CPG model however always uses the cosine of the phase (Equation (4)) to produce the output. Luckily, we can imitate a sine by a cosine by adding some phase bias to the cosine!:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba5967ae24fdc312",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "times = np.linspace(0, 2 * np.pi, 100)\n",
    "in_plane_positions = np.cos(times)\n",
    "out_of_plane_positions = np.cos(times - np.pi / 2)\n",
    "\n",
    "plt.scatter(in_plane_positions, out_of_plane_positions, c=times, cmap=\"Blues\")\n",
    "plt.xlabel(\"in-plane position\")\n",
    "plt.ylabel(\"out-of-plane position\")\n",
    "plt.show()\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "563d98d855ebc84",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Okay, so we know that we'll have to add a phase bias of $\\frac{\\pi}{2}$ between the IP and OOP oscillators.\n",
    "If we were to imitate rowing behaviour with our own arms, however, we'll notice that our left arm traverses the circle in the clockwise direction, while our right arm moves in the counterclockwise direction. The above motion goes in the counterclockwise direction. Luckily, to get a clockwise direction, we just need to play with the phase bias again. Instead of subtracting $\\frac{\\pi}{2}$ from the OOP phase, we just need to add it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2bda3025faa3405",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "times = np.linspace(0, 2 * np.pi, 100)\n",
    "in_plane_positions = np.cos(times)\n",
    "out_of_plane_positions = np.cos(times + np.pi / 2)\n",
    "\n",
    "plt.scatter(in_plane_positions, out_of_plane_positions, c=times, cmap=\"Blues\")\n",
    "plt.xlabel(\"in-plane position\")\n",
    "plt.ylabel(\"out-of-plane position\")\n",
    "plt.show()\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d06fb30206de66ec",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### CPG setup\n",
    "Now that we know how we can play with phase biases to produce circular motions in both the clockwise and counterclockwise directions, we can now implement our CPG system.\n",
    "To simplify things, we will not use a distinct oscillator for every joint in the brittle star's morphology. Instead, we'll just use two bi-directionally connected oscillators per arm: one for the IP motion and one for the OOP motion. We can then just map the output of an oscillator to all its respective joints in the arm (i.e. the IP oscillator's output is mapped to all IP joints of that arm, and the same thing for the OOP oscillator and the OOP joints). Given that we have 5 arms and 2 oscillators per arm, we thus have a total of $5 \\times 2 = 10$ oscillators. To make sure different arms stay synchronized, we will couple the IP oscillator of one arm to the IP oscillators of its neighbouring arms, and do the same thing for the OOP oscillators.\n",
    "\n",
    "![](https://github.com/Co-Evolve/SEL3-2024/blob/main/tutorials/assets/brittle_star_cpg.png?raw=true)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6aa3eaa9427864f8",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def create_cpg() -> CPG:\n",
    "    ip_oscillator_indices = jnp.arange(0, 10, 2)\n",
    "    oop_oscillator_indices = jnp.arange(1, 10, 2)\n",
    "\n",
    "    adjacency_matrix = jnp.zeros((10, 10))\n",
    "    # Connect oscillators within an arm\n",
    "    adjacency_matrix = adjacency_matrix.at[ip_oscillator_indices, oop_oscillator_indices].set(1)\n",
    "    # Connect IP oscillators of neighbouring arms\n",
    "    adjacency_matrix = adjacency_matrix.at[\n",
    "        ip_oscillator_indices, jnp.concatenate((ip_oscillator_indices[1:], jnp.array([ip_oscillator_indices[0]])))].set(\n",
    "            1\n",
    "            )\n",
    "    # Connect OOP oscillators of neighbouring arms\n",
    "    adjacency_matrix = adjacency_matrix.at[oop_oscillator_indices, jnp.concatenate(\n",
    "            (oop_oscillator_indices[1:], jnp.array([oop_oscillator_indices[0]]))\n",
    "            )].set(1)\n",
    "\n",
    "    # Make adjacency matrix symmetric (i.e. make all connections bi-directional)\n",
    "    adjacency_matrix = jnp.maximum(adjacency_matrix, adjacency_matrix.T)\n",
    "\n",
    "    return CPG(\n",
    "            weights=5 * adjacency_matrix,\n",
    "            amplitude_gain=20,\n",
    "            offset_gain=20,\n",
    "            dt=environment_configuration.control_timestep\n",
    "            )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d69e4f36160828b",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "A [brittle star's rowing gait](https://www.youtube.com/watch?v=X8UEST-flCM) divides the arms in roughly two roles: (1) one leading arm that points towards the direction of motion, and (2) rowers that pull the brittle star forwards. We will try to module our CPG in such a way that we can reproduce this behaviour.\n",
    "For the leading arm, this comes down to:\n",
    "- Setting its IP and OOP oscillator's amplitude $R$ to zero in order to inhibit oscillatory motion.\n",
    "    - Given that we inhibit the leading arms oscillatory motions, we can just set its coupling phase biases to zero as well.\n",
    "- Settings the offset $X$ of its IP oscillator to zero to make it point forwards.\n",
    "- Settings the offset $X$ if its OOP oscillator to a positive value to make it point upwards. This will avoid the leading arm sliding on the ground and creating unnecessary friction.\n",
    "\n",
    "For the rowers, it is a bit more complex:\n",
    "- We will set their amplitudes $R$ to the maximum joint range of motion.\n",
    "- We will set their IP and OOP offset $X$ to zero to make oscillations center around zero.\n",
    "- We will set their coupling phase biases $\\rho$ in such a way that their IP and OOP oscillators produce a circular motion of the arm. As discussed above, can be done by incorporating a phase bias of $\\frac{\\pi}{2}$ between the IP and OOP oscillator.\n",
    "    - Arms on the right of the leading arm will need to follow a counterclockwise direction, while arms on the left of the leading arm need to follow a clockwise direction.\n",
    "    - Let's say that for some rowing arm, our IP oscillator has index $i$ and our OOP oscillator has index $j$:\n",
    "        - To create a counterclockwise motion, the above test shows that we need $\\phi_{j} = \\phi_{i} - \\frac{\\pi}{2}$. This means that we'll need to set $\\rho_{ij} = -\\frac{\\pi}{2}$ and $\\rho_{ji} = \\frac{\\pi}{2}$. (If this confuses you, take another look at our very first CPG test and its corresponding phase plot.)\n",
    "        - To create a clockwise motion, the above test shows that we need $\\phi_{j} = \\phi_{i} + \\frac{\\pi}{2}$. This means that we'll need to set $\\rho_{ij} = \\frac{\\pi}{2}$ and $\\rho_{ji} = -\\frac{\\pi}{2}$.\n",
    "\n",
    "Since we connected the oscillators of neighbouring arms, we also need to double-check if these couplings require phase bias adaptations (as discussed in the explanation of the CPG model, we always want to ensure that all phase biases in a loop are consistent: they must sum to a multiple of $2\\pi$):\n",
    "- Phase biases of coupling's between the IP oscillators of the left rowers can be set to zero: we always want them to be completely in-sync.\n",
    "- Phase biases of coupling's between the IP oscillators of the right rowers can be set to zero: we always want them to be completely in-sync.\n",
    "- Phase biases of coupling's between the IP oscillators of a left and right rower can not be set to zero: they will need a phase bias of $\\pi$ as they were both shifted by $\\frac{\\pi}{2}$ in opposite directions due to the phase bias of the coupling with their respective OOP oscillators. In other words, without adaptation, this loop would have a phase bias sum of $\\frac{\\pi}{2} + \\frac{\\pi}{2} = \\pi$.\n",
    "- All the coupling's phase biases between OOP oscillators of neighbouring arms can be set to zero: we always want them to be completely in-sync as we for instance never want one arm going down while the others are going up.\n",
    "    - However, the coupling between oscillators of the leading arm and its neighbouring (rowing arms) presents an edge case. We have set the phase biases of the leading arm's oscillators to zero, while the neighbouring rowers will have non-zero phase biases between their IP and OOP oscillators. If we would just set the phase biases between the leading arm's oscillators and the neighbouring (rowing) arm's IP and OOP oscillators to zero, we get a mismatch! The leading arm's IP and OOP oscillators will try to be in-phase with each other, while also trying to be in-phase with the IP and OOP oscillators of the neighbouring arms which are out-of-phase. Luckily, as we set the amplitudes $R$ of our leading arm to zero, this mismatch would cause no harm. The change in phase frequency $\\dot\\phi_i$ that a coupling's phase bias $\\rho_{ij}$ can induce (Equation (1)) is scaled by the amplitude of the incoming oscillator $r_j$. Consequently, our rower's oscillators won't be influenced. Nevertheless, we'll protect ourselves from strange bugs that this issue might cause at a later stage (e.g. if we would use non-zero amplitudes for the leading arm) and fix this mismatch. Let's say that our leading arm's IP oscillator has index $x$ and its OOP oscillator has index $y$. As we have no phase bias between them, we need to make sure that $\\phi_{x} = \\phi_{y}$.\n",
    "        - For the first arm to the left of the leading arm (with IP oscillator index $i$ and OOP oscillator index $j$), we have set the phase biases such that we get $\\phi_{j} = \\phi_{i} + \\frac{\\pi}{2}$ to produce a clockwise circular motion.\n",
    "            - To fix the mismatch between the IP oscillators $x$ and $i$, we need to set their coupling's phase bias such that $\\phi_{x} = \\phi_{i} + \\frac{\\pi}{4}$. This means that we'll need to set $\\rho_{ix} = \\frac{\\pi}{4}$ and $\\rho_{xi} = -\\frac{\\pi}{4}$.\n",
    "            - To fix the mismatch between the OOP oscillators $y$ and $j$, we need to set their coupling's phase bias such that $\\phi_{y} = \\phi_{j} - \\frac{\\pi}{4}$. This means that we'll need to set $\\rho_{iy} = -\\frac{\\pi}{4}$ and $\\rho_{yi} = \\frac{\\pi}{4}$\n",
    "            - This works because we have $\\phi_{j} = \\phi_{i} + \\frac{\\pi}{2}$, which we can rewrite as $\\phi_{y} + \\frac{\\pi}{4} = \\phi_{x} - \\frac{\\pi}{4} + \\frac{\\pi}{2}$, which gives us $\\phi_y = \\phi_x$ \n",
    "        - Similarly, for the first arm to the right of the leading arm (with IP oscillator index $i$ and OOP oscillator index $j$), we have set the phase biases such that we get $\\phi_{j} = \\phi_{i} - \\frac{\\pi}{2}$ to produce a counterclockwise circular motion.\n",
    "            - To fix the mismatch between the IP oscillators $x$ and $i$, we need to set their coupling's phase bias such that $\\phi_{x} = \\phi_{i} - \\frac{\\pi}{4}$. This means that we'll need to set $\\rho_{ix} = -\\frac{\\pi}{4}$ and $\\rho_{xi} = \\frac{\\pi}{4}$\n",
    "            - To fix the mismatch between the OOP oscillators $y$ and $j$, we need to set their coupling's phase bias such that $\\phi_{y} = \\phi_{j} + \\frac{\\pi}{4}$. This means that we'll need to set $\\rho_{iy} = \\frac{\\pi}{4}$ and $\\rho_{yi} = -\\frac{\\pi}{4}$\n",
    "\n",
    "\n",
    "Note: It might be a bit tough to understand the reasoning behind the choices of phase biases and especially whether $\\rho_{ij}$ or $\\rho_{ji}$ should be negative. Whenever in doubt: play some more with the simple 2-oscillator model from above and look at the corresponding phase plot!   \n",
    "Note: The reason that we can come up with this kind of modulation that produces some desired gait is because we know the closed form solution of our CPG model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6dab680a71b92854",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "\n",
    "from typing import Tuple\n",
    "\n",
    "\n",
    "def get_oscillator_indices_for_arm(\n",
    "        arm_index: int\n",
    "        ) -> Tuple[int, int]:\n",
    "    return arm_index * 2, arm_index * 2 + 1\n",
    "\n",
    "\n",
    "@jax.jit\n",
    "def modulate_cpg(\n",
    "        cpg_state: CPGState,\n",
    "        leading_arm_index: int,\n",
    "        max_joint_limit: float\n",
    "        ) -> CPGState:\n",
    "    left_rower_arm_indices = [(leading_arm_index - 1) % 5, (leading_arm_index - 2) % 5]\n",
    "    right_rower_arm_indices = [(leading_arm_index + 1) % 5, (leading_arm_index + 2) % 5]\n",
    "\n",
    "    leading_arm_ip_oscillator_index, leading_arm_oop_oscillator_index = get_oscillator_indices_for_arm(\n",
    "            arm_index=leading_arm_index\n",
    "            )\n",
    "\n",
    "    R = jnp.zeros_like(cpg_state.R)\n",
    "    X = jnp.zeros_like(cpg_state.X)\n",
    "    rhos = jnp.zeros_like(cpg_state.rhos)\n",
    "    omegas = jnp.pi * jnp.ones_like(cpg_state.omegas)\n",
    "    phases_bias_pairs = []\n",
    "\n",
    "    def modulate_leading_arm(\n",
    "            _X: jnp.ndarray,\n",
    "            _arm_index: int\n",
    "            ) -> jnp.ndarray:\n",
    "        ip_oscillator_index, oop_oscillator_index = get_oscillator_indices_for_arm(arm_index=_arm_index)\n",
    "        return _X.at[oop_oscillator_index].set(max_joint_limit)\n",
    "\n",
    "    def modulate_left_rower(\n",
    "            _R: jnp.ndarray,\n",
    "            _arm_index: int\n",
    "            ) -> Tuple[jnp.ndarray, List[Tuple[int, int, float]]]:\n",
    "        ip_oscillator_index, oop_oscillator_index = get_oscillator_indices_for_arm(arm_index=_arm_index)\n",
    "        _R = _R.at[ip_oscillator_index].set(max_joint_limit)\n",
    "        _R = _R.at[oop_oscillator_index].set(max_joint_limit)\n",
    "        _phase_bias_pairs = [(ip_oscillator_index, oop_oscillator_index, jnp.pi / 2)]\n",
    "        return _R, _phase_bias_pairs\n",
    "\n",
    "    def phase_biases_first_left_rower(\n",
    "            _arm_index: int\n",
    "            ) -> List[Tuple[int, int, float]]:\n",
    "        ip_oscillator_index, oop_oscillator_index = get_oscillator_indices_for_arm(arm_index=_arm_index)\n",
    "        _phase_bias_pairs = [(ip_oscillator_index, leading_arm_ip_oscillator_index, jnp.pi / 4),\n",
    "                             (leading_arm_oop_oscillator_index, oop_oscillator_index, jnp.pi / 4)]\n",
    "        return _phase_bias_pairs\n",
    "\n",
    "    def modulate_right_rower(\n",
    "            _R: jnp.ndarray,\n",
    "            _arm_index: int\n",
    "            ) -> Tuple[jnp.ndarray, List[Tuple[int, int, float]]]:\n",
    "        ip_oscillator_index, oop_oscillator_index = get_oscillator_indices_for_arm(arm_index=_arm_index)\n",
    "        _R = _R.at[ip_oscillator_index].set(max_joint_limit)\n",
    "        _R = _R.at[oop_oscillator_index].set(max_joint_limit)\n",
    "        _phase_bias_pairs = [(oop_oscillator_index, ip_oscillator_index, jnp.pi / 2)]\n",
    "        return _R, _phase_bias_pairs\n",
    "\n",
    "    def phase_biases_first_right_rower(\n",
    "            _arm_index: int\n",
    "            ) -> List[Tuple[int, int, float]]:\n",
    "        ip_oscillator_index, oop_oscillator_index = get_oscillator_indices_for_arm(arm_index=_arm_index)\n",
    "        _phase_bias_pairs = [(leading_arm_ip_oscillator_index, ip_oscillator_index, jnp.pi / 4),\n",
    "                             (oop_oscillator_index, leading_arm_oop_oscillator_index, jnp.pi / 4)]\n",
    "        return _phase_bias_pairs\n",
    "\n",
    "    def phase_biases_second_rowers(\n",
    "            _left_arm_index: int,\n",
    "            _right_arm_index: int\n",
    "            ) -> List[Tuple[int, int, float]]:\n",
    "        left_ip_oscillator_index, _ = get_oscillator_indices_for_arm(arm_index=_left_arm_index)\n",
    "        right_ip_oscillator_index, _ = get_oscillator_indices_for_arm(arm_index=_right_arm_index)\n",
    "        _phase_bias_pairs = [(left_ip_oscillator_index, right_ip_oscillator_index, jnp.pi)]\n",
    "        return _phase_bias_pairs\n",
    "\n",
    "    X = modulate_leading_arm(_X=X, _arm_index=leading_arm_index)\n",
    "\n",
    "    R, phb = modulate_left_rower(_R=R, _arm_index=left_rower_arm_indices[0])\n",
    "    phases_bias_pairs += phb\n",
    "\n",
    "    R, phb = modulate_left_rower(_R=R, _arm_index=left_rower_arm_indices[1])\n",
    "    phases_bias_pairs += phb\n",
    "\n",
    "    R, phb = modulate_right_rower(_R=R, _arm_index=right_rower_arm_indices[0])\n",
    "    phases_bias_pairs += phb\n",
    "\n",
    "    R, phb = modulate_right_rower(_R=R, _arm_index=right_rower_arm_indices[1])\n",
    "    phases_bias_pairs += phb\n",
    "\n",
    "    phases_bias_pairs += phase_biases_first_left_rower(_arm_index=left_rower_arm_indices[0])\n",
    "    phases_bias_pairs += phase_biases_first_right_rower(_arm_index=right_rower_arm_indices[0])\n",
    "\n",
    "    phases_bias_pairs += phase_biases_second_rowers(\n",
    "            _left_arm_index=left_rower_arm_indices[1], _right_arm_index=right_rower_arm_indices[1]\n",
    "            )\n",
    "\n",
    "    for oscillator1, oscillator2, bias in phases_bias_pairs:\n",
    "        rhos = rhos.at[oscillator1, oscillator2].set(bias)\n",
    "        rhos = rhos.at[oscillator2, oscillator1].set(-bias)\n",
    "\n",
    "    # noinspection PyUnresolvedReferences\n",
    "    return cpg_state.replace(\n",
    "            R=R, X=X, rhos=rhos, omegas=omegas\n",
    "            )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a0c39530127caf2",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Before moving on, let's first do a validity check and plot our cpg outputs over some time:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15f5ab710f4fd64",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "cpg = create_cpg()\n",
    "cpg_state = cpg.reset(rng=jax.random.PRNGKey(seed=0))\n",
    "cpg_state = modulate_cpg(cpg_state=cpg_state, leading_arm_index=0, max_joint_limit=env.action_space.high[0] * 0.5)\n",
    "\n",
    "times = []\n",
    "outputs = []\n",
    "while cpg_state.time < environment_configuration.simulation_time:\n",
    "    outputs.append(cpg_state.outputs)\n",
    "    times.append(cpg_state.time)\n",
    "    cpg_state = cpg.step(\n",
    "            state=cpg_state\n",
    "            )\n",
    "\n",
    "import matplotlib.gridspec as gridspec\n",
    "\n",
    "outputs_per_oscillator = np.array(outputs).T\n",
    "gs = gridspec.GridSpec(int(cpg.num_oscillators), 1)\n",
    "fig = plt.figure()\n",
    "\n",
    "for i, outputs in enumerate(outputs_per_oscillator):\n",
    "    ax = fig.add_subplot(gs[i])\n",
    "    ax.plot(times, outputs)\n",
    "    ax.set_ylabel(f\"Oscillator {i}\", rotation=0, labelpad=20)\n",
    "    ax.set_ylim(env.action_space.low[0] - 0.1, env.action_space.high[1] + 0.1)\n",
    "    if i < outputs_per_oscillator.shape[0] - 1:\n",
    "        plt.tick_params(\n",
    "                axis='x', labelbottom='off'\n",
    "                )\n",
    "\n",
    "plt.xlabel(\"time\")\n",
    "plt.show()\n",
    "plt.show()\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8524d37cdaa8acd0",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Try to validate its correctness based on the rules that we implemented.\n",
    "\n",
    "Now we'll need to implement a helper function that maps our CPG's outputs to an actual action array. As stated above, we have two oscillators per arm, one for all in-plane motions within that arm and one for all out-of-plane motions for within that arm."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19b8936a0e380a87",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "@jax.jit\n",
    "def map_cpg_outputs_to_actions(\n",
    "        cpg_state: CPGState\n",
    "        ) -> jnp.ndarray:\n",
    "    num_arms = 5\n",
    "    num_oscillators_per_arm = 2\n",
    "    num_segments_per_arm = 5\n",
    "\n",
    "    cpg_outputs_per_arm = cpg_state.outputs.reshape((num_arms, num_oscillators_per_arm))\n",
    "    cpg_outputs_per_segment = cpg_outputs_per_arm.repeat(num_segments_per_arm, axis=0)\n",
    "\n",
    "    actions = cpg_outputs_per_segment.flatten()\n",
    "    return actions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b56e8d71a29c0881",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Now we can finally apply it to the actual brittle star:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ba47f9c758320e0",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "cpg = create_cpg()\n",
    "cpg_state = cpg.reset(rng=jax.random.PRNGKey(0))\n",
    "cpg_state = modulate_cpg(cpg_state=cpg_state, leading_arm_index=0, max_joint_limit=env.action_space.high[0] * 0.5)\n",
    "\n",
    "done = False\n",
    "frames = []\n",
    "env_state = jit_reset(rng=jax.random.PRNGKey(seed=0))\n",
    "while not (env_state.terminated | env_state.truncated):\n",
    "    cpg_state = cpg.step(state=cpg_state)\n",
    "    actions = map_cpg_outputs_to_actions(cpg_state=cpg_state)\n",
    "    env_state = jit_step(state=env_state, action=actions)\n",
    "    frame = post_render(env.render(state=env_state), environment_configuration=environment_configuration)\n",
    "    frames.append(frame)\n",
    "show_video(images=frames)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "194b38c3f1e38d53",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Looks good! Let's see if we can change the leading arm at half the simulation time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c90282842edbb36",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "cpg_state = cpg.reset(rng=jax.random.PRNGKey(0))\n",
    "cpg_state = modulate_cpg(cpg_state=cpg_state, leading_arm_index=0, max_joint_limit=env.action_space.high[0] * 0.5)\n",
    "\n",
    "done = False\n",
    "frames = []\n",
    "env_state = jit_reset(rng=jax.random.PRNGKey(seed=0))\n",
    "remodulated = False\n",
    "while not (env_state.terminated | env_state.truncated):\n",
    "    cpg_state = cpg.step(state=cpg_state)\n",
    "    actions = map_cpg_outputs_to_actions(cpg_state=cpg_state)\n",
    "    env_state = jit_step(state=env_state, action=actions)\n",
    "    frame = post_render(env.render(state=env_state), environment_configuration=environment_configuration)\n",
    "    frames.append(frame)\n",
    "\n",
    "    if not remodulated and env_state.info[\"time\"] > environment_configuration.simulation_time / 2:\n",
    "        cpg_state = modulate_cpg(\n",
    "                cpg_state=cpg_state, leading_arm_index=2, max_joint_limit=env.action_space.high[0] * 0.5\n",
    "                )\n",
    "        remodulated = True\n",
    "show_video(images=frames)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d161ce860f18dc16",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Awesome, by exploiting our prior knowledge on brittle star rowing behaviour and by making simplifications in our CPG system, we can manually define quite a decent gait for our brittle star robot! In applications where we cannot manually define the modulation parameters, we often apply optimization algorithms to solve this problem for us. For instance, we can optimize the modulation parameters to make the robot crawl as fast as possible or even to try to discover as many distinct gaits as possible. Optimization also allows us to scale up the number of oscillators (e.g. to the number of actuators). This gives the CPG more degrees of freedom to work with, allowing it to produce more complex behaviours than our current circular motions."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "605edfe8a326168c",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Excercises\n",
    "1. Now that we have found the modulation parameters that produce rowing behaviour, try to adapt the modulation parameters to produce reverse rowing behaviour.\n",
    "2. Try to apply the rowing gait to solve the targeted locomotion task. Every $n$ timesteps, you can do a check which arm is closest to the target. Then you can remodulate the CPG to make that arm the leading arm."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ed92a54193757dc",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Reading material\n",
    "- [Ijspeert, 2008, Central pattern generators for locomotion control in animals and robots: a review, Neural Networks](https://www.sciencedirect.com/science/article/pii/S0893608008000804)\n",
    "- [Sproewitz et al., 2008, Learning to move in modular robots using CPGs and online optimization, Int. J. Robotics Res.](https://journals.sagepub.com/doi/10.1177/0278364907088401)\n",
    "- [Buchli et al., 2006, Engineering entrainment and adaptation in limit cycle systems, Biol. Cybern.](https://pubmed.ncbi.nlm.nih.gov/17146662/)\n",
    "- [Ijspeert et al., 2007, From swimming to walking with a salamander robot driven by a spinal cord model, Science](https://www.science.org/doi/10.1126/science.1138353)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "faf963b81d4b2bc0",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "example-env-kernel",
   "language": "python",
   "name": "example-env-kernel"
  },
  "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.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}