tutorials/brb_brittle_star_qlearning_tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d19319df3ee69e93",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "# <h1><center>Q-learning Tutorial</center></h1>\n",
    "\n",
    "This notebook provides an introductory tutorial to Q-learning. Specifically, we will implement tabular Q-Learning in JAX and use it to steer a simplified brittle star robot toward a target location. The [brittle star robot and its environment](https://github.com/Co-Evolve/brt/tree/main/biorobot/brittle_star) is part of the [**the Bio-inspired Robotics Testbed (BRT)**](https://github.com/Co-Evolve/brt). Instead of directly outputting joint-level motor commands, we will use our Q-learning controller to select motion primitives from a predefined set. These motion primitives were the final result of the CPG tutorial and are simplified rowing gaits that differentiate based on the leading arm. In essence, our Q-learning controller will thus learn how to map the brittle star's state to the correct leading arm to steer it toward a target location. \n",
    "\n",
    "The main goal of this tutorial is to give you a practical example that you can play around with to gain intuition in the fundamental design choices of reinforcement learning experiments (state representations, action definitions, reward definitions, etc.)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32ecd84ea867e0d5",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Q-Learning\n",
    "\n",
    "Q-learning belongs to the class of [model-free](https://en.wikipedia.org/wiki/Model-free_(reinforcement_learning)) reinforcement algorithms. This means that the algorithm does not require prior knowledge (i.e. a model) of the environment. It also belongs to the class of off-policy algorithms, meaning that it does not use the 'current policy' to produce actions at every timestep.\n",
    "\n",
    "As its name entails, the algorithm's goal is to learn the function $Q(s, a)$, which represents the expected cumulative reward for taking a particular action $a$ in a given state $s$. In other words, we try to learn a function that predicts the expected payoff of doing a certain action in a given state. Given such a function, we can select the action with the highest expected payoff in every state.\n",
    "\n",
    "In this tutorial, we will focus on tabular Q-learning. This means that our $Q(s, a)$ function is implemented as a table with discretized states as rows and discretized actions as columns. After a successful optimization, each cell in this table will approximate the Q-value of a state and an action. To get a good action given some state, we can then just select the action that corresponds to the highest Q-value of the row indexed by the discretized state.\n",
    "\n",
    "Initially, this Q-Table will be populated with zeros or small random values (small random values giving the additional benefit of encouraging some extra initial exploration). The Q-Learning algorithm then tries to optimize the values in this table through an iterative process of exploration and exploitation. During the exploration phase, the agent will take random actions to get information about the environment and update the Q-Table accordingly. As the agent explores more, it gradually transitions into the exploitation phase, where it leverages the learned Q-values to make more informed decisions and maximize the cumulative reward further.\n",
    "\n",
    "The tabular Q-learning algorithm can be summarized as follows:\n",
    "1. Initialize the Q-table with arbitrary values or zeros.\n",
    "2. Observe the current state $s$ of the environment.\n",
    "3. Choose an action $a$ to take based on an exploration-exploitation trade-off. This can for instance be done by using the epsilon-greedy approach, where the agent selects a random action with probability $\\epsilon$ (exploration) and chooses the action with the highest Q-value with a complementary probability $(1-\\epsilon)$ (exploitation).\n",
    "4. Perform the chosen action $a$, transition to the next state $s'$ and observe the reward $r$.\n",
    "5. Update the Q-Table of the state-action pair using the Q-learning update rule<br>\n",
    "        $Q(s,a) = (1 - \\alpha) Q(s,a) + \\alpha(r + \\gamma\\max_{a' \\in A}(Q(s', a')))$,<br>\n",
    "        where alpha $\\alpha$ is the learning rate, gamma $\\gamma$ is the discount factor that determines the importance of future rewards, $r$ is the immediate reward obtained (given by the environment), and $\\max_{a' \\in A}(Q(s', a'))$ represents the maximum Q-value for the next state\n",
    "6. Repeat steps $2$ to $5$ until convergence or a predefined number of iterations.\n",
    "\n",
    "A natural extension of the tabular Q-learning algorithm (which is limited to discrete actions and states) is the [Deep Q-Learning (DQN) algorithm](https://huggingface.co/learn/deep-rl-course/unit3/introduction). As its name gives away, the DQN algorithm swaps the Q-table for a deep neural network, which enables mapping continuous states to $Q(s, a)$ values (the network will have one output neuron per action)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1594f329c5fed9fb",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Implementing tabular Q-learning in JAX\n",
    "\n",
    "When implementing something in JAX it's important to remember that JAX follows the [functional programming paradigm](https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#pure-functions). Put simply, we thus rely on pure functions (deterministic and without side effects) and immutable data structures (instead of changing data in place, new data structures are created with the desired modifications) as primary building blocks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8bb1f509d9b18a0",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import functools\n",
    "from typing import Tuple\n",
    "\n",
    "import chex\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "from flax import struct\n",
    "\n",
    "\n",
    "@struct.dataclass\n",
    "class QLearningPolicyParameters:\n",
    "    q_table: jnp.ndarray\n",
    "    alpha: float\n",
    "    epsilon: float\n",
    "    gamma: float\n",
    "\n",
    "\n",
    "class QLearningPolicy:\n",
    "    def __init__(\n",
    "            self,\n",
    "            num_states: int,\n",
    "            num_actions: int\n",
    "            ) -> None:\n",
    "        self._num_states = num_states\n",
    "        self._num_actions = num_actions\n",
    "\n",
    "    @functools.partial(jax.jit, static_argnums=(0,))\n",
    "    def apply_q_learning_update_rule(\n",
    "            self,\n",
    "            policy_parameters: QLearningPolicyParameters,\n",
    "            state_index: int,\n",
    "            next_state_index: int,\n",
    "            action_index: int,\n",
    "            reward: float\n",
    "            ) -> QLearningPolicyParameters:\n",
    "        old_q_value = policy_parameters.q_table[state_index, action_index]\n",
    "        best_future_q_value = jnp.max(policy_parameters.q_table[next_state_index])\n",
    "        q_value_update = reward + policy_parameters.gamma * best_future_q_value\n",
    "        new_q_value = (1 - policy_parameters.alpha) * old_q_value + policy_parameters.alpha * q_value_update\n",
    "\n",
    "        new_q_table = policy_parameters.q_table.at[state_index, action_index].set(new_q_value)\n",
    "        # noinspection PyUnresolvedReferences\n",
    "        return policy_parameters.replace(q_table=new_q_table)\n",
    "\n",
    "    @functools.partial(jax.jit, static_argnums=(0,))\n",
    "    def epsilon_greedy_action(\n",
    "            self,\n",
    "            policy_parameters: QLearningPolicyParameters,\n",
    "            state_index: int,\n",
    "            rng: chex.PRNGKey\n",
    "            ) -> Tuple[QLearningPolicyParameters, int]:\n",
    "        explore_rng, random_action_rng = jax.random.split(rng, 2)\n",
    "        explore = jax.random.uniform(explore_rng) < policy_parameters.epsilon\n",
    "\n",
    "        def get_random_action() -> int:\n",
    "            return jax.random.choice(key=random_action_rng, a=jnp.arange(policy_parameters.q_table.shape[1]))\n",
    "\n",
    "        def get_greedy_action() -> int:\n",
    "            return jnp.argmax(policy_parameters.q_table[state_index])\n",
    "\n",
    "        action_index = jax.lax.cond(\n",
    "                pred=explore, true_fun=get_random_action, false_fun=get_greedy_action\n",
    "                )\n",
    "\n",
    "        return action_index\n",
    "\n",
    "    @functools.partial(jax.jit, static_argnums=(0,))\n",
    "    def reset(\n",
    "            self,\n",
    "            rng: chex.PRNGKey,\n",
    "            alpha: float,\n",
    "            gamma: float,\n",
    "            epsilon: float\n",
    "            ) -> QLearningPolicyParameters:\n",
    "        # noinspection PyArgumentList\n",
    "        return QLearningPolicyParameters(\n",
    "                q_table=jax.random.uniform(\n",
    "                        key=rng,\n",
    "                        shape=(self._num_states, self._num_actions),\n",
    "                        dtype=jnp.float32,\n",
    "                        minval=-0.001,\n",
    "                        maxval=0.001\n",
    "                        ), alpha=alpha, epsilon=epsilon, gamma=gamma\n",
    "                )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4f3b6a2ead0ed7",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Case study: Selecting behavioural primitives for directed brittle star locomotion\n",
    "\n",
    "Now that we have implemented our `QLearningPolicy`, we can apply it to learn how to steer our brittle star toward a target location. To simplify the problem, we will keep the target location fixed. Furthermore, instead of directly outputting joint-level actions, we will use our Q-learning controller to select behavioural primitives out of a predefined set.\n",
    "These behavioural primitives were the final result of the CPG tutorial: simplified rowing gaits that differentiate based on the leading arm. In essence, our Q-learning controller will thus learn how to map states to a leading arm. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "acf71a68021b859b",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### Environment setup\n",
    "\n",
    "Let's start by setting up our brittle star environment. We will use the directed locomotion variant (which we will always supply with the same `jax.random.PRNGKey` in order to keep the target location fixed).\n",
    "The following cell will first do some preliminary checks to make sure the underlying physics engine (MuJoCo) is correctly loaded and to make sure that JAX can access the GPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3df4a5004f1aa81",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import subprocess\n",
    "import logging\n",
    "import numpy as np\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",
    "    import jax.numpy as jnp\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",
    "np.set_printoptions(precision=3, suppress=True, linewidth=100)\n",
    "\n",
    "jnp.set_printoptions(precision=3, suppress=True, linewidth=100)\n",
    "\n",
    "print('MuJoCo installation successful.')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7cf1378d-0570-43b9-9efe-ebe27c4b80af",
   "metadata": {},
   "source": [
    "This next cell (similar to previous tutorials) defines the `morphology_specification` (i.e. the brittle star morphology), the `arena_configuration` (i.e. some settings w.r.t. the aquarium in which we place the brittle star) and the `environment_configuration` (which defines and configures the directed locomotion task). The cell also implements some utility functions for visualization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9c86f67b8e0018d",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from biorobot.brittle_star.environment.directed_locomotion.shared import \\\n",
    "    BrittleStarDirectedLocomotionEnvironmentConfiguration\n",
    "import numpy as np\n",
    "from moojoco.environment.base import MuJoCoEnvironmentConfiguration\n",
    "from biorobot.brittle_star.environment.directed_locomotion.dual import BrittleStarDirectedLocomotionEnvironment\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=3, use_p_control=True, use_torque_control=False\n",
    "        )\n",
    "arena_configuration = AquariumArenaConfiguration(\n",
    "        size=(1.5, 1.5), sand_ground_color=False, attach_target=True, wall_height=1.5, wall_thickness=0.1\n",
    "        )\n",
    "environment_configuration = BrittleStarDirectedLocomotionEnvironmentConfiguration(\n",
    "        target_distance=1.5,\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() -> BrittleStarDirectedLocomotionEnvironment:\n",
    "    morphology = MJCFBrittleStarMorphology(\n",
    "            specification=morphology_specification\n",
    "            )\n",
    "    arena = MJCFAquariumArena(\n",
    "            configuration=arena_configuration\n",
    "            )\n",
    "    env = BrittleStarDirectedLocomotionEnvironment.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",
    "        path: str | None = 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",
    "    if path:\n",
    "        media.write_video(path=path, images=filtered_images)\n",
    "    return media.show_video(images=filtered_images)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b8266e5-8d3f-46f4-ac9c-b9a647396683",
   "metadata": {},
   "source": [
    "As always, we `jax.jit` the two main environment functions: `step` and `reset`. As stated above, we will always pass the same Pseudo Random Number Generator (PRNG) key to the reset call in order to keep the target location fixed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51b4b9cc27fe20f6",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from functools import partial\n",
    "\n",
    "rng = jax.random.PRNGKey(seed=0)\n",
    "env = create_environment()\n",
    "\n",
    "# Fix the target location -> always use the same seed\n",
    "env_fixed_target_reset_fn = jax.jit(partial(env.reset, rng=jax.random.PRNGKey(seed=12)))\n",
    "env_step_fn = jax.jit(env.step)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7aa3ceae-fe6a-48df-b6c9-e5618399f353",
   "metadata": {},
   "source": [
    "The next cell prints out the environment's observation space, action space and the `info` dictionary that our environment updates every step. As you can see, we explicitly reset the environment with seed $12$, which causes the target to appear in the upper right corner."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "738cb21ac5a5650e",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "rng, sub_rng = jax.random.split(rng, 2)\n",
    "env_state = env_fixed_target_reset_fn()\n",
    "print(\"Observation space:\")\n",
    "print(env.observation_space)\n",
    "print()\n",
    "print(\"Action space:\")\n",
    "print(env.action_space)\n",
    "print()\n",
    "print(\"Info:\")\n",
    "print(env_state.info)\n",
    "print()\n",
    "media.show_image(post_render(env.render(env_state), environment_configuration=env.environment_configuration))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58cd1a24-5d87-4424-9337-a6c3a649555b",
   "metadata": {},
   "source": [
    "### Experimental design\n",
    "\n",
    "Given our brittle star environment and task definition, we can now design our reinforcement learning setup. The three main components are the state representation, the actions, and the reward function. When experimenting with reinforcement learning, it is incredibly important to always maintain a very clear overview of these three components (and of course the environment/morphology configuration as well). When the learning fails or is slow, it is almost always one of these design choices that lies at the root of it (given that there are no bugs on the environment side).\n",
    "\n",
    "\n",
    "#### State representation\n",
    "\n",
    "One of the most important design choices when applying a Q-learning algorithm (and reinforcement learning algorithms in general), is how to represent the state. In other words, what kind of information do we supply to the policy to make it return good actions? Which observations that the environment returns do we use and do we transform them in some way to represent them better? When deciding on this, a good exercise is to imagine that you are the robot trying to solve the given task; Ask yourself what kind of information you would require. In this case, we also need a discretized state representation, since we're working with tabular Q-Learning.\n",
    "\n",
    "To keep things as simple as possible, in this tutorial, we will approach this locomotion task from a grid-based perspective. We will discretize the 2D aquarium floor in a 6x6 grid of square cells (one cell of this grid is equal to a small square that the rendered aquarium floor's texture above indicates). As a discrete state representation for this locomotion task with a fixed target location, we will take the index of the grid cell that the brittle star robot's disk is currently in. Take a moment to think about this state representation. If you were the robot, would this allow you to solve the task? \n",
    "\n",
    "Sidenote: During this project, we will always be using simulation environments. Simulation environments are nice since they give us full control over and full observability of everything in the environment (for instance the exact position of the brittle star's disk). In the real world, however, we do not get this for free. There we have to rely purely on sensory information (e.g. cameras, touch sensors, motor encoders...) and additional processing to get this kind of information (which is often noisy). Always keep this in mind, and even though we won't do a sim2real transfer during this project, always be able to explain how you would allow your selected state representation to be implemented on a real-world robot. For instance, in this grid-based representation, we could use a downwards-facing camera above the aquarium. Using some computer vision techniques, the brittle star's disk position can be extracted.\n",
    "\n",
    "This next cell implements our grid-based state representation. We name the function 'state_indexer' as it extracts a row index for our Q-Table."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9529b201-f39c-4e3f-bcfc-a06eb1716e5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from moojoco.environment.mjx_env import MJXEnvState\n",
    "\n",
    "ARENA_SIZE = jnp.array(arena_configuration.size)\n",
    "NUM_CELLS_PER_AXIS = 6\n",
    "NUM_POSITION_STATES = NUM_CELLS_PER_AXIS ** 2\n",
    "TOTAL_NUM_STATES = NUM_POSITION_STATES\n",
    "\n",
    "\n",
    "def position_to_state_index(\n",
    "        position: jnp.ndarray\n",
    "        ) -> jnp.ndarray:\n",
    "    # position is in [-arena size, arena size], first convert it to [0, 2 * arena size]\n",
    "    shifted_position = position + ARENA_SIZE\n",
    "    normalized_position = shifted_position / (2 * ARENA_SIZE + 0.001)\n",
    "\n",
    "    x, y = (NUM_CELLS_PER_AXIS * normalized_position).astype(jnp.int32)\n",
    "    return x + NUM_CELLS_PER_AXIS * y\n",
    "\n",
    "\n",
    "@jax.jit\n",
    "def state_indexer(\n",
    "        env_state: MJXEnvState\n",
    "        ) -> int:\n",
    "    robot_position = env_state.observations[\"disk_position\"][:2]\n",
    "    position_index = position_to_state_index(position=robot_position)\n",
    "    return position_index"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2150d8bc-c7a7-4755-8d8d-0bf2ee29d465",
   "metadata": {},
   "source": [
    "Quick test to check the correctness of the position to grid cell index:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2d054be-5f97-4b17-89d0-f6042d7496a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "x_positions = jnp.linspace(-ARENA_SIZE[0], ARENA_SIZE[0], NUM_CELLS_PER_AXIS)\n",
    "y_positions = jnp.linspace(-ARENA_SIZE[1], ARENA_SIZE[1], NUM_CELLS_PER_AXIS)\n",
    "positions = jnp.array(jnp.meshgrid(x_positions, y_positions)).T.reshape(-1, 2)\n",
    "\n",
    "positions = positions[jnp.argsort(positions[:, 1], kind='stable'), :]\n",
    "for position in positions:\n",
    "    print(f\"Position: {position}\\t->\\tState index: {position_to_state_index(position)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b604a397298308ee",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "#### Actions\n",
    "Our action selects one of the behavioral primitives that we converged to at the end of the CPG tutorial: a rowing gait that is defined by a leading arm. In other words, we define an action here as the index of the leading arm.\n",
    "\n",
    "The next cell first copies the CPG implementation, the CPG creation, the leading-arm-based modulation, and the CPG readout function (i.e. CPG state to joint-level actuator actions) from the CPG tutorial."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c2214b36e742650",
   "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",
    "from functools import partial\n",
    "from typing import Tuple\n",
    "\n",
    "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\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, ) -> None:\n",
    "        self._weights = weights\n",
    "        self._amplitude_gain = amplitude_gain\n",
    "        self._offset_gain = offset_gain\n",
    "        self._dt = dt\n",
    "        self._solver = euler_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",
    "        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",
    "                )\n",
    "\n",
    "\n",
    "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",
    "            )\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",
    "        joint_limit: float, ) -> CPGState:\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 = 2 * jnp.pi * jnp.ones_like(cpg_state.omegas)\n",
    "    phases_bias_pairs = []\n",
    "\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",
    "    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(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(joint_limit)\n",
    "        _R = _R.at[oop_oscillator_index].set(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(joint_limit)\n",
    "        _R = _R.at[oop_oscillator_index].set(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",
    "    cpg_state = cpg_state.replace(\n",
    "            R=R, X=X, rhos=rhos, omegas=omegas\n",
    "            )\n",
    "\n",
    "    return cpg_state\n",
    "\n",
    "\n",
    "@jax.jit\n",
    "def map_cpg_outputs_to_actions(\n",
    "        cpg_state: CPGState\n",
    "        ) -> jnp.ndarray:\n",
    "    num_arms = morphology_specification.number_of_arms\n",
    "    num_oscillators_per_arm = 2\n",
    "    num_segments_per_arm = morphology_specification.number_of_segments_per_arm[0]\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": "25f769cd0bed53df",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Given the CPG implementation above, we can now implement our 'action mapper', i.e. the function that maps a discrete action index (column in our Q-Table, which in this case indicates the leading arm) to the actuator-level control commands our environment expects.  To recap, our entire pipeline from environment observations to actuator controls will look like this:\n",
    "\n",
    "<br>`environment observations` $\\rightarrow$ `state index (denoting the grid cell the robot's central disk is in)` $\\rightarrow$ `Epsilon-greedy policy` $\\rightarrow$ `action index (a leading arm index)` $\\rightarrow$ `CPG modulation` $\\rightarrow$ `CPG outputs` $\\rightarrow$ `actuator controls`<be>\n",
    "\n",
    "Here we introduce two additional simplifications of the learning task. \n",
    "\n",
    "The first simplification is that we do not modulate the CPG every control timestep. Given the environment configuration, the current control timestep is $0.04$ seconds (you can extract this using `environment_configuration.control_timestep`, the cell above also sets this as the time increment for our CPG's differential equation solver). This means that without any adaptations, we would re-modulate our CPG every $0.04$ seconds. Although this might be beneficial in a more complex environment, e.g. one in which there are obstacles that require faster adaptation, in this case, it makes the learning more difficult. During the exploration phase, the controller will output a lot of random actions. Consequently, if we re-modulate the CPG with a random action this frequently, we will switch between leading arms quite often and the robot will barely move. Therefore, when the policy proposes a leading arm index, we will always take $50$ environment steps with this same leading arm.\n",
    "\n",
    "The second simplification is that we will try to make subsequent actions independent of one another. If we do not do this and purely rely on the CPG to smoothly transition between two different actions (i.e. two different leading arms, and thus two different CPG states), we often get some minor erratic movements. While this does not pose the greatest of challenges for a well-trained controller, it makes learning more difficult and we want to keep things as simple as possible here. We will do this by always warming up the CPG first (i.e. forwarding it a couple of steps with zero amplitudes) and finishing the action generation with a cool down (i.e. setting the amplitudes back to zero a bit before the end) with a cooldown. The warm up will allow the imposed phase biases between oscillators to be reached. The cool-down will force the robot to return to a fixed position (all joint positions at $0$) after it rowed with the given leading arm."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c50f3682bd2b5e60",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Dict\n",
    "\n",
    "NUM_ACTIONS = 5\n",
    "\n",
    "NUM_WARMUP_STEPS = 5\n",
    "NUM_COOLDOWN_STEPS = 5\n",
    "NUM_SUBSTEPS_PER_MODULATION = 50\n",
    "NUM_ROWING_STEPS = NUM_SUBSTEPS_PER_MODULATION - NUM_WARMUP_STEPS - NUM_COOLDOWN_STEPS\n",
    "\n",
    "\n",
    "@functools.partial(jax.jit, static_argnums=(0,))\n",
    "def cpg_action_mapper(\n",
    "        cpg: CPG,\n",
    "        action_index: int,\n",
    "        joint_limit: float\n",
    "        ) -> jnp.ndarray:\n",
    "    # We'll just always reset the CPG to the original state\n",
    "    cpg_state = cpg.reset(jax.random.PRNGKey(0))\n",
    "\n",
    "    modulated_cpg_state = modulate_cpg(cpg_state=cpg_state, leading_arm_index=action_index, joint_limit=joint_limit)\n",
    "\n",
    "    def _step_cpg(\n",
    "            _cpg_state,\n",
    "            _\n",
    "            ):\n",
    "        _cpg_state = cpg.step(state=_cpg_state)\n",
    "        _actions = map_cpg_outputs_to_actions(cpg_state=_cpg_state)\n",
    "        return _cpg_state, _actions\n",
    "\n",
    "    # Warm up\n",
    "    cpg_state = modulated_cpg_state.replace(R=modulated_cpg_state.R * 0)\n",
    "    cpg_state, warmup_actions = jax.lax.scan(_step_cpg, cpg_state, (), NUM_WARMUP_STEPS)\n",
    "\n",
    "    # Rowing\n",
    "    cpg_state = cpg_state.replace(R=modulated_cpg_state.R)\n",
    "    cpg_state, rowing_actions = jax.lax.scan(_step_cpg, cpg_state, (), NUM_ROWING_STEPS)\n",
    "\n",
    "    # Cool down\n",
    "    cpg_state = cpg_state.replace(R=modulated_cpg_state.R * 0)\n",
    "    cpg_state, cooldown_actions = jax.lax.scan(_step_cpg, cpg_state, (), NUM_COOLDOWN_STEPS)\n",
    "\n",
    "    all_actions = jnp.concatenate((warmup_actions, rowing_actions, cooldown_actions), axis=0)\n",
    "\n",
    "    return all_actions\n",
    "\n",
    "\n",
    "def _step_env_n_times(\n",
    "        env_state: MJXEnvState,\n",
    "        actions: jnp.ndarray\n",
    "        ) -> MJXEnvState:\n",
    "    def _step_env(\n",
    "            _env_state,\n",
    "            action\n",
    "            ):\n",
    "        _env_state = env_step_fn(_env_state, action)\n",
    "        return _env_state, _env_state\n",
    "\n",
    "    final_env_state, env_states = jax.lax.scan(\n",
    "            _step_env, env_state, actions\n",
    "            )\n",
    "\n",
    "    env_state = final_env_state.replace(\n",
    "            reward=jnp.max(env_states.reward),\n",
    "            terminated=jnp.any(env_states.terminated),\n",
    "            truncated=jnp.any(env_states.truncated)\n",
    "            )\n",
    "\n",
    "    return env_state\n",
    "\n",
    "\n",
    "def _step_env_n_times_with_frames(\n",
    "        env_state: MJXEnvState,\n",
    "        actions: jnp.ndarray\n",
    "        ) -> Tuple[MJXEnvState, List[np.ndarray]]:\n",
    "    frames = []\n",
    "    for action in actions:\n",
    "        next_env_state = env_step_fn(env_state, action)\n",
    "        env_state = next_env_state.replace(\n",
    "                reward=max(next_env_state.reward, env_state.reward),\n",
    "                terminated=next_env_state.terminated | env_state.terminated,\n",
    "                truncated=next_env_state.truncated | env_state.truncated\n",
    "                )\n",
    "        frame = post_render(env.render(state=env_state), environment_configuration=environment_configuration)\n",
    "        frames.append(frame)\n",
    "    return env_state, frames"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8fcb8863797ad3d6",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "##### Visual validation\n",
    "\n",
    "Let's first check the correctness of our current code (mainly the action mapper) by visualizing some episodes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8be555d161f0df77",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def visualize_episode(\n",
    "        policy_parameters: QLearningPolicyParameters,\n",
    "        rng: chex.PRNGKey,\n",
    "        time_limit: float | None = None\n",
    "        ) -> None:\n",
    "    env_state = env_fixed_target_reset_fn()\n",
    "\n",
    "    # noinspection PyUnresolvedReferences\n",
    "    # make the policy deterministic\n",
    "    policy_parameters = policy_parameters.replace(epsilon=0)\n",
    "\n",
    "    frames = []\n",
    "\n",
    "    while (time_limit is None or env_state.info[\"time\"] < time_limit) and not (\n",
    "            env_state.terminated | env_state.truncated):\n",
    "        state_index = state_indexer(env_state)\n",
    "        action_index = q_learning_policy.epsilon_greedy_action(\n",
    "                policy_parameters=policy_parameters, state_index=state_index, rng=rng\n",
    "                )\n",
    "        actions = cpg_action_mapper(\n",
    "                cpg=cpg, action_index=action_index, joint_limit=env.action_space.high[0] * 0.5\n",
    "                )\n",
    "        env_state, _frames = _step_env_n_times_with_frames(env_state, actions)\n",
    "        frames += _frames\n",
    "\n",
    "    show_video(images=frames)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "071b3b0c-268d-4702-8531-a006687d4ef0",
   "metadata": {},
   "source": [
    "For each action (i.e. leading arm index), we will run one episode in which we always use that action. This can be done by overwriting the underlying Q-Table and setting a high value for that action's column. If we then set the epsilon parameter to zero, the `QLearningPolicy` will always select the desired action."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18293f997dd8edb5",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "q_learning_policy = QLearningPolicy(num_states=TOTAL_NUM_STATES, num_actions=NUM_ACTIONS)\n",
    "cpg = create_cpg()\n",
    "\n",
    "rng, episode_rng, q_learner_rng = jax.random.split(rng, 3)\n",
    "policy_parameters = q_learning_policy.reset(rng=q_learner_rng, alpha=0.01, gamma=0.90, epsilon=0.0)\n",
    "\n",
    "for action_index, label in enumerate([\"first arm\", \"second arm\", \"third arm\", \"fourth arm\", \"fifth arm\"]):\n",
    "    adapted_policy_parameters = policy_parameters.replace(q_table=policy_parameters.q_table.at[:, action_index].set(10))\n",
    "    print(f\"Action: {label}\")\n",
    "    visualize_episode(policy_parameters=adapted_policy_parameters, rng=episode_rng, time_limit=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff8c64e9-9789-4836-81ac-6bfff5345f5b",
   "metadata": {},
   "source": [
    "Looking good! As you can notice, however, even with our additional simplifications, the brittle star does not perfectly follow the direction of the leading arm. This makes our state transition slightly more stochastic. Nevertheless, if we collect enough data, our Q-Learning algorithm should be able to handle this complexity."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e597640-95ac-4386-b82a-acbbddb53cad",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-05T15:16:09.469975Z",
     "start_time": "2024-02-05T15:16:09.446350Z"
    }
   },
   "source": [
    "#### Reward\n",
    "The last important design choice when applying RL is the reward function. The reward is our main way to tell our policy what we want it to learn.\n",
    "The `BrittleStarDirectedLocomotionEnvironment` defines the reward as the difference between the previous distance to the target and the current distance to the target ([implementation](https://github.com/Co-Evolve/brb/blob/new-framework/brb/brittle_star/environment/directed_locomotion/shared.py#L28)). While this is a solid default choice, here we will use a slightly different reward $r$:\n",
    "\n",
    "$r = \n",
    "\\begin{cases} \n",
    "+20 & \\text{if target reached} \\\\\n",
    "-0.001 & \\text{otherwise}\n",
    "\\end{cases}$\n",
    "\n",
    "This reward provides less 'information' than the original reward; it is a ['sparse' reward](https://www.geeksforgeeks.org/sparse-rewards-in-reinforcement-learning/). We will use this reward because it gives a nice example of how important exploration is; If the agent does not explore well enough, it will never reach the target and thus never receive the positive reward. Consequently, it will never learn. The negative constant reward that is given at every timestep can be seen as a time penalty; I.e. the longer it takes for the agent to reach the goal, the smaller the total reward. This will stimulate our policy to reach the target as fast as possible.\n",
    "\n",
    "Note: keep in mind that you're never limited to the reward function that the environment provides by default! If you want to change something environment-related, take a look at how the environment implements it and modify it however you like (the necessary `mujoco.MjModel`, `mujoco.MjData`, `mjx.Model`, and `mjx.Data` are always returned in the `MJXEnvState`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88013d34-e00c-49d3-85c8-e4a4a738ce1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from moojoco.environment.mjx_env import MJXEnvState\n",
    "\n",
    "\n",
    "def calculate_sparse_reward(\n",
    "        state: MJXEnvState\n",
    "        ) -> float:\n",
    "    pred = state.terminated  # If this is True, we have reached the target\n",
    "\n",
    "    def time_penalty() -> float:\n",
    "        return -0.001\n",
    "\n",
    "    def target_reached_reward() -> float:\n",
    "        return 20.0\n",
    "\n",
    "    return jax.lax.cond(pred, target_reached_reward, time_penalty)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "269ebc34458d3b8f",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### Rollout function\n",
    "Now that we have implemented our experimental setup, we can now implement our 'rollout' or 'playout' function.\n",
    "This function runs an entire simulation episode, using the policy defined by given parameters. Normally, a rollout function just returns the collected data (i.e. $(s, a, s', r)$ samples. In this case, however, we want our policy to be updated after every step (since we are using temporal difference learning). Consequently, here we will also return the updated policy parameters.\n",
    "\n",
    "Note: The rollout function implementation in the next cell does not take early environment terminations into account (these occur when the robot has reached the target location before the allocated time). The reason for this is that we use the [jax.lax.scan](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.scan.html) function, which always runs a given function for a fixed number of times. While we could also use a [jax.lax.while_loop](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.while_loop.html) to allow early stopping, the general rule-of-thumb when using JAX effectively is to run loops for a fixed number of times. So instead, the implementation below just stops updating the policy parameters once the episode is terminated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9141041fce69178c",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from typing import Union\n",
    "\n",
    "\n",
    "def rollout_with_policy_updates(\n",
    "        rng: chex.PRNGKey,\n",
    "        policy_parameters: QLearningPolicyParameters,\n",
    "        reward_calculator: Callable[[MJXEnvState], float]\n",
    "        ) -> Tuple[QLearningPolicyParameters, Dict[str, jnp.ndarray]]:\n",
    "    \"\"\"\n",
    "    Do a single episode rollout and update the policy every step.\n",
    "    \"\"\"\n",
    "    env_state = env_fixed_target_reset_fn()\n",
    "\n",
    "    def policy_step(\n",
    "            _state: Tuple[MJXEnvState, QLearningPolicyParameters, chex.PRNGKey, bool],\n",
    "            _: None\n",
    "            ) -> Tuple[Tuple[MJXEnvState, QLearningPolicyParameters, chex.PRNGKey, bool], Dict[str, Union[int, float]]]:\n",
    "        # Do environment step\n",
    "        _env_state, _policy_parameters, _policy_rng, _episode_done = _state\n",
    "        _state_index = state_indexer(_env_state)\n",
    "\n",
    "        _policy_rng, _sub_rng = jax.random.split(_policy_rng, 2)\n",
    "        _action_index = q_learning_policy.epsilon_greedy_action(\n",
    "                policy_parameters=_policy_parameters, state_index=_state_index, rng=_sub_rng\n",
    "                )\n",
    "        _actions = cpg_action_mapper(\n",
    "                cpg=cpg, action_index=_action_index, joint_limit=env.action_space.high[0] * 0.5\n",
    "                )\n",
    "        _next_env_state = _step_env_n_times(env_state=_env_state, actions=_actions)\n",
    "        _next_state_index = state_indexer(_next_env_state)\n",
    "\n",
    "        # Update policy\n",
    "        def _fake_update() -> QLearningPolicyParameters:\n",
    "            return _policy_parameters\n",
    "\n",
    "        def _real_update() -> QLearningPolicyParameters:\n",
    "            _updated_policy_parameters = q_learning_policy.apply_q_learning_update_rule(\n",
    "                    policy_parameters=_policy_parameters,\n",
    "                    state_index=_state_index,\n",
    "                    next_state_index=_next_state_index,\n",
    "                    action_index=_action_index,\n",
    "                    reward=reward_calculator(_next_env_state)\n",
    "                    )\n",
    "            return _updated_policy_parameters\n",
    "\n",
    "        _next_policy_parameters = jax.lax.cond(_episode_done, _fake_update, _real_update)\n",
    "\n",
    "        _episode_done |= _next_env_state.terminated | _next_env_state.truncated\n",
    "\n",
    "        carry = (_next_env_state, _next_policy_parameters, _policy_rng, _episode_done)\n",
    "        return carry, {\n",
    "                \"reward\": reward_calculator(_next_env_state), \"done\": _episode_done,\n",
    "                \"target_reached\": _next_env_state.terminated, \"state_index\": _state_index,\n",
    "                \"action_index\": _action_index}\n",
    "\n",
    "    (_, policy_parameters, _, _), scan_out = jax.lax.scan(\n",
    "            policy_step,\n",
    "            (env_state, policy_parameters, rng, False),\n",
    "            (),\n",
    "            env.environment_configuration.total_num_control_steps / NUM_SUBSTEPS_PER_MODULATION\n",
    "            )\n",
    "\n",
    "    return policy_parameters, scan_out"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8aa7502b-8e96-415c-8b9d-2685345e15ab",
   "metadata": {},
   "source": [
    "### Logging via WandB\n",
    "We will be using [Weights And Biases (W&B)](https://docs.wandb.ai/quickstart) for logging. W&B is an easy-to-use experiment tracker, and allows us to log to an online dashboard (hosted on the cloud) directly from our python code. Before continuing, checkout the following [quickstart](https://docs.wandb.ai/quickstart) and create an account."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e33987d-473f-4ed2-8b90-cc6265203470",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "!{sys.executable} -m pip install -q wandb\n",
    "import wandb"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b089b00-8747-4747-9d98-76cb2f99d527",
   "metadata": {},
   "source": [
    "Given some rollout data, let's first implement a utility function that extracts some basic but interesting metrics to log."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b638b0b-7b81-470c-b965-88de3d763684",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_logging_metrics(\n",
    "        policy_parameters: QLearningPolicyParameters,\n",
    "        rollout_data: Dict[str, jnp.ndarray],\n",
    "        total_num_episodes: int\n",
    "        ) -> Dict[str, Union[int, float, wandb.Histogram]]:\n",
    "    done_timestep = jnp.argmax(rollout_data[\"done\"]) + 1\n",
    "\n",
    "    metrics = {\n",
    "            \"Task/cumulative-reward\": jnp.sum(rollout_data[\"reward\"][:done_timestep]),\n",
    "            \"Task/succes\": int(jnp.any(rollout_data[\"target_reached\"])),\n",
    "            \"Q-Table/average-value\": jnp.average(policy_parameters.q_table),\n",
    "            \"Hyperparameters/epsilon\": policy_parameters.epsilon, \"General/total-num-episodes\": total_num_episodes}\n",
    "\n",
    "    states = rollout_data[\"state_index\"]\n",
    "    max_bin = TOTAL_NUM_STATES\n",
    "    if TOTAL_NUM_STATES > 512:\n",
    "        reduction_factor = TOTAL_NUM_STATES / 512\n",
    "        states = states / reduction_factor\n",
    "        max_bin = 512\n",
    "\n",
    "    visited_states_histogram = np.histogram(\n",
    "            a=states, bins=max_bin, range=(0, max_bin)\n",
    "            )\n",
    "    metrics[\"Task/visited-states\"] = wandb.Histogram(np_histogram=visited_states_histogram)\n",
    "    selected_actions_histogram = np.histogram(\n",
    "            a=rollout_data[\"action_index\"], bins=NUM_ACTIONS, range=(0, NUM_ACTIONS)\n",
    "            )\n",
    "    metrics[\"Task/selected-actions\"] = wandb.Histogram(np_histogram=selected_actions_histogram)\n",
    "\n",
    "    return metrics"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c41caa8f27ff1093",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### Our first learning experiment\n",
    "\n",
    "Now we truly have everything to do our very first learning experiment! The following cell starts by initializing our parameters, and then iteratively runs a rollout with policy updates. It also includes some minimal logging to W&B to track the progress.\n",
    "Note: the very first iteration might take a bit longer as this also includes the jit compilation of the `rollout_with_policy_updates` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3a824cb41ecaf87",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "\n",
    "NUM_EPISODES = 100\n",
    "\n",
    "rng = jax.random.PRNGKey(seed=0)\n",
    "rng, q_learner_rng = jax.random.split(rng, 2)\n",
    "policy_parameters = q_learning_policy.reset(rng=q_learner_rng, alpha=0.01, gamma=0.95, epsilon=0.25)\n",
    "\n",
    "jitted_rollout_with_policy_updates = jax.jit(\n",
    "        partial(rollout_with_policy_updates, reward_calculator=calculate_sparse_reward)\n",
    "        )\n",
    "\n",
    "wandb.init(\n",
    "        project=\"SEL3-2024-QL-Tutorial\", config={\n",
    "                \"alpha\": policy_parameters.alpha, \"gamma\": policy_parameters.gamma,\n",
    "                \"epsilon\": policy_parameters.epsilon}\n",
    "        )\n",
    "\n",
    "for episode_i in tqdm(range(NUM_EPISODES), desc=\"Training policy\"):\n",
    "    rng, sub_rng = jax.random.split(rng, 2)\n",
    "    policy_parameters, rollout_data = jitted_rollout_with_policy_updates(\n",
    "            rng=sub_rng, policy_parameters=policy_parameters\n",
    "            )\n",
    "\n",
    "    metrics = get_logging_metrics(\n",
    "            policy_parameters=policy_parameters, rollout_data=rollout_data, total_num_episodes=episode_i\n",
    "            )\n",
    "    wandb.log(metrics)\n",
    "\n",
    "wandb.finish()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f875ba83-9f9b-4f7e-9725-38d9f216e8c7",
   "metadata": {},
   "source": [
    "### Result analysis\n",
    "\n",
    "After running an experiment, always take some time to analyze the resulting plots on W&B. Are the results what you expected them to be? Next to analyzing the resulting data, we often also take an empirical look at what kind of behavior the learned controller produces:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7e25dfd-b8fd-4a40-a0ff-8ad79fa2d22c",
   "metadata": {},
   "outputs": [],
   "source": [
    "rng, episode_rng = jax.random.split(rng, 2)\n",
    "visualize_episode(policy_parameters=policy_parameters, rng=episode_rng)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3dd6d6fc-8633-47fb-a102-98588e7627e7",
   "metadata": {},
   "source": [
    "As you can see, this controller does not work! The robot does not reach the target location. The reason for this, as indicated by the plots shown on W&B, is that we never reach the target location during training. This means that the policy never receives any reward other than the time penalty and thus never receives a usable reward signal. Simply said, it just can't learn what to do because it's never told it did something good.\n",
    "\n",
    "### Enriching the reward signal\n",
    "One possible solution to this issue is to provide a richer reward signal, i.e. a reward function that gives our policy more information on how well its doing. One way to do so is to enrich the time penalty reward with the original environment's reward: the amount of distance the robot got closer to the target in this timestep compared to the previous timestep. If the robot then starts moving away from the target location, it will get penalized more and consequently, this reward provides a direct stimulation to move toward the target. So instead of relying on the exploration mechanism to randomly stumble on the positive reward, we let each state transition provide some more information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d60ca8d4-f078-493e-b13a-1cf9054edced",
   "metadata": {},
   "outputs": [],
   "source": [
    "from moojoco.environment.mjx_env import MJXEnvState\n",
    "\n",
    "\n",
    "def calculate_rich_reward(\n",
    "        state: MJXEnvState\n",
    "        ) -> float:\n",
    "    pred = state.terminated  # If this is True, we have reached the target\n",
    "\n",
    "    def distance_delta_reward() -> float:\n",
    "        # We can weight the importance of components of the reward\n",
    "        return 2 * state.reward - 0.001\n",
    "\n",
    "    def target_reached_reward() -> float:\n",
    "        return 20.0\n",
    "\n",
    "    return jax.lax.cond(pred, target_reached_reward, distance_delta_reward)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ebb2d8c-1704-4f05-ada7-1c0fa4367b4e",
   "metadata": {},
   "source": [
    "Let's redo the training with this updated reward function!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51d1d54d-4fa3-4882-9bb2-f68fceda1843",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "\n",
    "NUM_EPISODES = 100\n",
    "\n",
    "rng = jax.random.PRNGKey(seed=0)\n",
    "rng, q_learner_rng = jax.random.split(rng, 2)\n",
    "policy_parameters = q_learning_policy.reset(rng=q_learner_rng, alpha=0.01, gamma=0.95, epsilon=0.25)\n",
    "\n",
    "jitted_rollout_with_policy_updates = jax.jit(\n",
    "        partial(rollout_with_policy_updates, reward_calculator=calculate_rich_reward)\n",
    "        )\n",
    "\n",
    "wandb.init(\n",
    "        project=\"SEL3-2024-QL-Tutorial\", config={\n",
    "                \"alpha\": policy_parameters.alpha, \"gamma\": policy_parameters.gamma,\n",
    "                \"epsilon\": policy_parameters.epsilon}\n",
    "        )\n",
    "\n",
    "for episode_i in tqdm(range(NUM_EPISODES), desc=\"Training policy\"):\n",
    "    rng, sub_rng = jax.random.split(rng, 2)\n",
    "    policy_parameters, rollout_data = jitted_rollout_with_policy_updates(\n",
    "            rng=sub_rng, policy_parameters=policy_parameters\n",
    "            )\n",
    "\n",
    "    metrics = get_logging_metrics(\n",
    "            policy_parameters=policy_parameters, rollout_data=rollout_data, total_num_episodes=episode_i\n",
    "            )\n",
    "    wandb.log(metrics)\n",
    "\n",
    "wandb.finish()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48ca0e35-5d08-445b-a53d-7fce25732e1d",
   "metadata": {},
   "source": [
    "The plots on W&B look more promising now! Let's visualize the behaviour:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9bdc58b2-22a6-44ff-ab88-60b06a1bc135",
   "metadata": {},
   "outputs": [],
   "source": [
    "rng, episode_rng = jax.random.split(rng, 2)\n",
    "visualize_episode(policy_parameters=policy_parameters, rng=episode_rng)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eed06c31-907a-4ac1-8dcc-470e048a395d",
   "metadata": {},
   "source": [
    "Although the controller now thus gets quite close to the target, in the training plots, it does not seem to consistently reach the target. The main reason for this is that we keep our epsilon value fixed at $25\\%$ during the entire training time. This means that $25\\%$ of all actions are randomly selected. So even though the policy starts learning which action is actually optimal, we keep forcing it to do random actions as well. While this randomness is required at the beginning of training (because the policy does not know what to do), we want to reduce its influence when the policy gets better.\n",
    "\n",
    "#### Rebalancing the exploration-exploitation trade-off over time: epsilon decay\n",
    "We can solve this issue by incorporating epsilon decay. In other words, we will reduce our epsilon value over time. This will allow the policy to start with a high exploration factor, and gradually convert this into a high exploitation factor (i.e. exploiting the learned knowledge)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e66f54bd-ce1d-4d90-ac30-8ca02454f93b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def epsilon_decay(\n",
    "        policy_parameters: QLearningPolicyParameters,\n",
    "        start: float,\n",
    "        end: float,\n",
    "        progress: float\n",
    "        ) -> QLearningPolicyParameters:\n",
    "    return policy_parameters.replace(epsilon=start + progress * (end - start))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4bd00dd0-ed08-4e1a-9326-12b4e062c6a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "NUM_EPISODES = 100\n",
    "\n",
    "rng = jax.random.PRNGKey(seed=0)\n",
    "rng, q_learner_rng = jax.random.split(rng, 2)\n",
    "policy_parameters = q_learning_policy.reset(rng=q_learner_rng, alpha=0.01, gamma=0.95, epsilon=0.25)\n",
    "\n",
    "jitted_rollout_with_policy_updates = jax.jit(\n",
    "        partial(rollout_with_policy_updates, reward_calculator=calculate_rich_reward)\n",
    "        )\n",
    "\n",
    "wandb.init(\n",
    "        project=\"SEL3-2024-QL-Tutorial\",\n",
    "        config={\n",
    "                \"alpha\": policy_parameters.alpha, \"gamma\": policy_parameters.gamma,\n",
    "                \"epsilon\": policy_parameters.epsilon}\n",
    "        )\n",
    "\n",
    "for episode_i in tqdm(range(NUM_EPISODES), desc=\"Training policy\"):\n",
    "    rng, sub_rng = jax.random.split(rng, 2)\n",
    "\n",
    "    progress = jnp.clip(episode_i / 75, 0, 1)\n",
    "    policy_parameters = epsilon_decay(policy_parameters=policy_parameters, start=0.8, end=0.0, progress=progress)\n",
    "\n",
    "    policy_parameters, rollout_data = jitted_rollout_with_policy_updates(\n",
    "            rng=sub_rng, policy_parameters=policy_parameters\n",
    "            )\n",
    "\n",
    "    metrics = get_logging_metrics(\n",
    "            policy_parameters=policy_parameters, rollout_data=rollout_data, total_num_episodes=episode_i\n",
    "            )\n",
    "    wandb.log(metrics)\n",
    "\n",
    "wandb.finish()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4b680a0-2384-4b62-9082-fa3459c90469",
   "metadata": {},
   "outputs": [],
   "source": [
    "rng, episode_rng = jax.random.split(rng, 2)\n",
    "visualize_episode(policy_parameters=policy_parameters, rng=episode_rng)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f5484bb-e060-4b0b-bfe9-6376e7298d93",
   "metadata": {},
   "source": [
    "Seems to be working! The policy now consistently reaches the target location at the end of training."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "affd42fbc860baa6",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### Exploit JAX: vectorize!\n",
    "\n",
    "As always, since we're running things on the GPU, it would be unwise to not exploit its parallelization capabilities. Generally, there are two ways in which we can exploit vectorization in RL: (1) We can train multiple distinct agents in parallel (e.g. when we want to compare different hyperparameter settings or if we want to check if our method is stable with respect to stochasticity), or (2) We can train a single agent and parallelize over the rollouts (e.g. to increase data collection speed and increase exploration due to the increase in different trajectories).\n",
    "\n",
    "In this case, we will go for option two and train one single agent using multiple rollouts. In this case, our rollout function also does policy updates. This means that if we vectorize it through `jax.vmap`, it will return a batch of different `QLearningPolicyParameters` with different versions of the Q-Table which we'll need to merge. In this case, we'll do this merge by averaging over the different versions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91765082-bae0-476a-96eb-46668048703b",
   "metadata": {},
   "outputs": [],
   "source": [
    "vectorized_rollout_with_updates = jax.jit(\n",
    "        jax.vmap(partial(rollout_with_policy_updates, reward_calculator=calculate_rich_reward), in_axes=(0, None), )\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b08bc183-1706-4c81-82ea-e05a35f96ddd",
   "metadata": {},
   "source": [
    "We'll have to rewrite our `get_logging_metrics` function such that it can handle batched rollout data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92296bb8-f74e-48ad-8b8c-00304222996e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Union\n",
    "\n",
    "\n",
    "def get_vectorized_logging_metrics(\n",
    "        policy_parameters: QLearningPolicyParameters,\n",
    "        rollout_data: Dict[str, jnp.ndarray],\n",
    "        total_num_episodes: int\n",
    "        ) -> Dict[str, Union[int, float, wandb.Histogram]]:\n",
    "    num_parallel_envs = rollout_data[\"done\"].shape[0]\n",
    "\n",
    "    done_timestep = jnp.argmax(rollout_data[\"done\"], axis=1)\n",
    "\n",
    "    metrics = {\n",
    "            \"Task/average-cumulative-reward\": jnp.average(\n",
    "                    jnp.cumsum(rollout_data[\"reward\"], axis=1)[jnp.arange(num_parallel_envs), done_timestep]\n",
    "                    ), \"Task/average-succes\": jnp.average(jnp.any(rollout_data[\"target_reached\"], axis=1)),\n",
    "            \"Q-Table/average-value\": jnp.average(policy_parameters.q_table),\n",
    "            \"Hyperparameters/epsilon\": policy_parameters.epsilon, \"General/total-num-episodes\": total_num_episodes}\n",
    "\n",
    "    visited_states_histogram = np.histogram(\n",
    "            a=rollout_data[\"state_index\"], bins=TOTAL_NUM_STATES, range=(0, TOTAL_NUM_STATES)\n",
    "            )\n",
    "    metrics[\"Task/visited-states\"] = wandb.Histogram(np_histogram=visited_states_histogram)\n",
    "    selected_actions_histogram = np.histogram(\n",
    "            a=rollout_data[\"action_index\"], bins=NUM_ACTIONS, range=(0, NUM_ACTIONS)\n",
    "            )\n",
    "    metrics[\"Task/selected-actions\"] = wandb.Histogram(np_histogram=selected_actions_histogram)\n",
    "\n",
    "    return metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8354cf29-c3b8-47a0-a3bd-03e5bc076471",
   "metadata": {},
   "outputs": [],
   "source": [
    "NUM_EPISODES = 100\n",
    "NUM_PARALLEL_ENVS = 32\n",
    "\n",
    "rng = jax.random.PRNGKey(seed=0)\n",
    "rng, q_learner_rng = jax.random.split(rng, 2)\n",
    "policy_parameters = q_learning_policy.reset(rng=q_learner_rng, alpha=0.01, gamma=0.95, epsilon=0.25)\n",
    "\n",
    "wandb.init(\n",
    "        project=\"SEL3-2024-QL-Tutorial\", config={\n",
    "                \"alpha\": policy_parameters.alpha, \"gamma\": policy_parameters.gamma,\n",
    "                \"epsilon\": policy_parameters.epsilon, \"num_parallel_envs\": NUM_PARALLEL_ENVS}\n",
    "        )\n",
    "\n",
    "for episode_i in tqdm(range(NUM_EPISODES), desc=\"Training policy\"):\n",
    "    rng, *sub_rngs = jax.random.split(rng, NUM_PARALLEL_ENVS + 1)\n",
    "\n",
    "    progress = jnp.clip(episode_i / 75, 0, 1)\n",
    "    policy_parameters = epsilon_decay(policy_parameters=policy_parameters, start=0.8, end=0.0, progress=progress)\n",
    "\n",
    "    batch_policy_parameters, rollout_data = vectorized_rollout_with_updates(jnp.array(sub_rngs), policy_parameters)\n",
    "\n",
    "    policy_parameters = batch_policy_parameters.replace(\n",
    "            q_table=jnp.average(batch_policy_parameters.q_table, axis=0),\n",
    "            epsilon=batch_policy_parameters.epsilon[0],\n",
    "            alpha=batch_policy_parameters.alpha[0],\n",
    "            gamma=batch_policy_parameters.gamma[0]\n",
    "            )\n",
    "\n",
    "    metrics = get_vectorized_logging_metrics(\n",
    "            policy_parameters=policy_parameters,\n",
    "            rollout_data=rollout_data,\n",
    "            total_num_episodes=episode_i * NUM_PARALLEL_ENVS\n",
    "            )\n",
    "\n",
    "    wandb.log(metrics)\n",
    "\n",
    "wandb.finish()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a6bb946-abab-431e-b184-beb022ba621a",
   "metadata": {},
   "source": [
    "Nice! Simply by vectorizing over 32 environments, we have just (approximately) increased our data collection speed from 100 episodes in 5.5 minutes ($\\pm18$ episodes per minute) to 3200 episodes in 12 minutes ($\\pm266$ episodes per minute)!\n",
    "Time to visualize again!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de6e5cad-6e5f-4552-93f3-d2c5f1df1070",
   "metadata": {},
   "outputs": [],
   "source": [
    "rng, episode_rng = jax.random.split(rng, 2)\n",
    "visualize_episode(policy_parameters=policy_parameters, rng=episode_rng)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b391090-70b6-4161-a55b-99a43989e72c",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "\n",
    "In this tutorial we have implemented and applied a simple reinforcement learning technique, tabular Q-Learning, to steer a brittle star robot towards a fixed target location. Instead of directly outputting joint-level motor commands, we selected motion primitives from an archive of leading-arm-based gaits. \n",
    "\n",
    "Experimenting with this setup has shown us the importance of good state, action, and reward signal design. Reinforcement learning algorithms are commonly very fragile w.r.t. these design choices and their hyperparameters. Therefore, it's very important to always have a clear overview of all design choices made (this is incredibly helpful for you when debugging why something is not learning and necessary for others to be able to help you with it).\n",
    "\n",
    "While we already have quite a nice controller, this controller still lacks a lot of adaptability. What if for instance, the target is at a random position? What if the floor is not flat or the aquarium contains obstacles? What if we lose an arm? Throughout this project, you will iteratively try to increase the controller's adaptiveness. This will most often complexify the task that the controller must be able to do, requiring you to improve your reinforcement learning setup. At a certain point, this will require you to shift toward more advanced algorithms such as [Deep Q-Learning](https://arxiv.org/abs/1312.5602) or [Proximal Policy Optimization](https://arxiv.org/abs/1707.06347).\n",
    "\n",
    "In essence, the end goal of the reinforcement learning track in this project will be to directly modulate the CPG based on given observations to steer the brittle star away from light (i.e. the light escape environment variant will be used). This will be done by either (1) Directly outputting the modulation parameters (amplitudes, offsets, frequencies, phase biases) of an underlying CPG system or (2) Selecting from a broader set of motion primitives (generated by the Quality-Diversity track)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "745532fe57357b2c",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Some recommended next steps\n",
    "\n",
    "In general: try to improve the learning as well as possible so that you can complexify the task. Always try to reason about and predict the influence of a certain modification before training, and compare your predictions with the actual results afterward! This is the best and fastest way to **improve your intuition, which is the main goal of this tutorial**!\n",
    "\n",
    "* Improve the logging\n",
    "    * Visualize an evaluation episode every $n$ rollout (with epsilon set to 0) and log the video on W&B. This will allow you to empirically analyze how the behavior changes over time (handy for debugging, and interesting to see).\n",
    "* Try to improve the training\n",
    "    * We currently set the hyperparameters (alpha, epsilon, gamma) ad hoc. Use JAX's vectorization capabilities to tune this.\n",
    "    * Test different reward functions.\n",
    "    * Introduce variation in the target location.\n",
    "        * You will need to supply the `env.reset` call with a different seed every time. \n",
    "        * You will need to adapt the state as this currently does not represent the target location.\n",
    "        * You can try to define a curriculum (start with an env that has its targets closer, then gradually increase the distance to the target).\n",
    "    * Make the state representation position independent (i.e. move away from the grid perspective). Position-dependent states (especially when working with a table) are often not the best way to represent the state, as there are a lot of options!\n",
    "* Try to make the controller decentralized (i.e. instead of one controller that outputs an action for each arm, develop a controller architecture that has one controller per arm)\n",
    "    * The brittle star's nervous system is decentralized as well. Developing a decentralized controller will thus bring you one step closer to the biological system, making your controller slightly more 'bio-inspired'.\n",
    "    * You can share the controller's parameters between different arms (every arm must be able to do what every other arm does anyway, they just need to output different actions  based on the different states they are in). You will thus need to make the states arm-specific.\n",
    "    * One possible approach:\n",
    "        * Arm-specific action: defines the role of the arm (leading arm, left rower, or right rower)\n",
    "        * Arm-specific state: angle between (1) the direction that the arm is pointing towards and (2) the direction towards the target\n",
    "        * You'll have to modify the CPG and its modulation function as well!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "048510bf-3c4c-42d6-a152-26f90d1f8d03",
   "metadata": {},
   "source": [
    "## Recommended reading material\n",
    "\n",
    "* [How to train your robot with deep reinforcement learning: lessons we have learned](https://arxiv.org/abs/2102.02915)\n",
    "* [Playing Atari with Deep Reinforcement Learning](https://arxiv.org/abs/1312.5602)\n",
    "* [Proximal Policy Optimization Algorithms](https://arxiv.org/abs/1707.06347)\n",
    "* [Deep Reinforcement Learning that Matters](https://arxiv.org/abs/1709.06560)\n",
    "* [CPG-RL: Learning Central Pattern Generators for Quadruped Locomotion](https://arxiv.org/abs/2211.00458)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d85e7ea6-04b6-457d-8478-0d4bc0ddf2b6",
   "metadata": {},
   "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
}