{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6icfEHC8fe6T"
      },
      "outputs": [],
      "source": [
        "from dataclasses import dataclass\n",
        "import numpy as np\n",
        "import mediapy as media\n",
        "from pathlib import Path\n",
        "import enum\n",
        "from tqdm import tqdm\n",
        "import mujoco"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AA6YeD4BfrJB"
      },
      "source": [
        "## Helper methods"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9JzEeCBCfodw"
      },
      "outputs": [],
      "source": [
        "class Resolution(enum.Enum):\n",
        "  SD = (480, 640)\n",
        "  HD = (720, 1280)\n",
        "  UHD = (2160, 3840)\n",
        "\n",
        "\n",
        "def quartic(t: float) -\u003e float:\n",
        "  return 0 if abs(t) \u003e 1 else (1 - t**2) ** 2\n",
        "\n",
        "\n",
        "def blend_coef(t: float, duration: float, std: float) -\u003e float:\n",
        "  normalised_time = 2 * t / duration - 1\n",
        "  return quartic(normalised_time / std)\n",
        "\n",
        "\n",
        "def unit_smooth(normalised_time: float) -\u003e float:\n",
        "  return 1 - np.cos(normalised_time * 2 * np.pi)\n",
        "\n",
        "\n",
        "def azimuth(\n",
        "    time: float, duration: float, total_rotation: float, offset: float\n",
        ") -\u003e float:\n",
        "  return offset + unit_smooth(time / duration) * total_rotation"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "shIh7-9DftEO"
      },
      "source": [
        "## Parameters"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "X0RdpXB4fuPu"
      },
      "outputs": [],
      "source": [
        "res = Resolution.SD\n",
        "fps = 60\n",
        "duration = 10.0\n",
        "ctrl_rate = 2\n",
        "ctrl_std = 0.05\n",
        "total_rot = 60\n",
        "blend_std = .8"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NPwLUhFOfvKF"
      },
      "outputs": [],
      "source": [
        "h, w = res.value"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_sBeUwo9fw0M"
      },
      "source": [
        "## Loading and rendering the model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PCat4NfbfytO"
      },
      "outputs": [],
      "source": [
        "model_dir = Path(\"unitree_go1\")\n",
        "model_xml = model_dir / \"scene.xml\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "q3asoEdrfzip"
      },
      "outputs": [],
      "source": [
        "# Load model.\n",
        "model = mujoco.MjModel.from_xml_path(str(model_xml))\n",
        "data = mujoco.MjData(model)\n",
        "\n",
        "# Make sure offscreen rendering can support the desired resolution.\n",
        "model.vis.global_.offheight = h\n",
        "model.vis.global_.offwidth = w\n",
        "\n",
        "renderer = mujoco.Renderer(model, height=h, width=w)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dZXxPKcJf15T"
      },
      "outputs": [],
      "source": [
        "mujoco.mj_forward(model, data)\n",
        "renderer.update_scene(data)\n",
        "media.show_image(renderer.render())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FWga8gT2f3cq"
      },
      "source": [
        "## Checking for keyframes"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NwNYSexff2re"
      },
      "outputs": [],
      "source": [
        "for key in range(model.nkey):\n",
        "  mujoco.mj_resetDataKeyframe(model, data, key)\n",
        "  mujoco.mj_forward(model, data)\n",
        "  renderer.update_scene(data)\n",
        "  media.show_image(renderer.render())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0kwngycGf64m"
      },
      "source": [
        "## Render!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zIGr56xUf7gR"
      },
      "outputs": [],
      "source": [
        "np.random.seed(12345)\n",
        "\n",
        "# Rendering options for visual and collision geoms.\n",
        "vis = mujoco.MjvOption()\n",
        "vis.geomgroup[2] = True\n",
        "vis.geomgroup[3] = False\n",
        "coll = mujoco.MjvOption()\n",
        "coll.geomgroup[2] = False\n",
        "coll.geomgroup[3] = True\n",
        "coll.flags[mujoco.mjtVisFlag.mjVIS_CONVEXHULL] = True\n",
        "\n",
        "# Create a camera that will revolve around the robot.\n",
        "camera = mujoco.MjvCamera()\n",
        "mujoco.mjv_defaultFreeCamera(model, camera)\n",
        "camera.distance = 1\n",
        "offset = model.vis.global_.azimuth\n",
        "\n",
        "# Sample actuator noise and smooth it.\n",
        "nsteps = int(np.ceil(duration / model.opt.timestep))\n",
        "perturb = np.random.randn(nsteps, model.nu)\n",
        "width = int(nsteps * ctrl_rate / duration)\n",
        "kernel = np.exp(-0.5 * np.linspace(-3, 3, width) ** 2)\n",
        "kernel /= np.linalg.norm(kernel)\n",
        "for i in range(model.nu):\n",
        "  perturb[:, i] = np.convolve(perturb[:, i], kernel, mode=\"same\")\n",
        "\n",
        "# Set the desired control point.\n",
        "if model.nkey \u003e 0:\n",
        "  mujoco.mj_resetDataKeyframe(model, data, 0)\n",
        "  ctrl0 = data.ctrl.copy()\n",
        "else:\n",
        "  mujoco.mj_resetData(model, data)\n",
        "  ctrl0 = np.mean(model.actuator_ctrlrange, axis=1)\n",
        "\n",
        "frames = []\n",
        "for i in tqdm(range(nsteps)):\n",
        "  data.ctrl[:] = ctrl0 + ctrl_std * perturb[i]\n",
        "  mujoco.mj_step(model, data)\n",
        "  if len(frames) \u003c data.time * fps:\n",
        "    camera.azimuth = azimuth(data.time, duration, total_rot, offset)\n",
        "    renderer.update_scene(data, camera, scene_option=vis)\n",
        "    vispix = renderer.render().copy().astype(np.float32)\n",
        "    renderer.update_scene(data, camera, scene_option=coll)\n",
        "    collpix = renderer.render().copy().astype(np.float32)\n",
        "    b = blend_coef(data.time, duration, blend_std)\n",
        "    frame = (1 - b) * vispix + b * collpix\n",
        "    frame = frame.astype(np.uint8)\n",
        "    frames.append(frame)\n",
        "media.show_video(frames, fps=fps, loop=False)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Tmw1thGfhA7B"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "colab": {
      "private_outputs": true,
      "provenance": [
        {
          "file_id": "1Pr1Dc7KRgt4h4i26Y1Knf6hxeffPDVDV",
          "timestamp": 1696017196949
        }
      ]
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}