{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MpkYHwCqk7W-"
   },
   "source": [
    "![MuJoCo banner](https://raw.githubusercontent.com/google-deepmind/mujoco/main/banner.png)\n",
    "\n",
    "## Учим робота ходить — плейбук мастер-класса\n",
    "\n",
    "Институт информационных технологий РТУ МИРЭА\n",
    "\n",
    "«Как обучить робота ходить, видеть, говорить и… думать?»\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xBSdkbmGN2K-"
   },
   "source": [
    "### Copyright notice"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_UbO9uhtBSX5"
   },
   "source": [
    "> <p><small><small>Copyright 2026 Google LLC.</small></p>\n",
    "> <p><small><small>Licensed under the Apache License, Version 2.0 (the \"License\"); you may not use this file except in compliance with the License. You may obtain a copy of the License at <a href=\"https://www.apache.org/licenses/LICENSE-2.0\">https://www.apache.org/licenses/LICENSE-2.0</a>.</small></small></p>\n",
    "> <p><small><small>Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License.</small></small></p>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dNIJkb_FM2Ux"
   },
   "source": [
    "# Плейбук: учим робота ходить в симуляторе\n",
    "\n",
    "Это учебный ноутбук из проекта **MuJoCo Playground** (Google DeepMind), переведённый для нашего мастер-класса. Мы шаг за шагом обучим виртуального четвероногого робота **Unitree Go1** ходить — тем же методом, которым учат настоящую Панду.\n",
    "\n",
    "**Для обучения обязательно нужен GPU-рантайм Colab.** Проверить и включить: меню **«Среда выполнения» → «Сменить среду выполнения»** → выбрать **GPU** (подойдёт бесплатный T4).\n",
    "\n",
    "Оригинальный ноутбук: [google-deepmind/mujoco_playground](https://github.com/google-deepmind/mujoco_playground/blob/main/learning/notebooks/locomotion.ipynb) (лицензия Apache 2.0, уведомление об авторских правах — ниже).\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "workshop-intro"
   },
   "source": [
    "## Что мы сделаем (карта ноутбука)\n",
    "\n",
    "1. **Подготовка** — поставим библиотеки и проверим, что GPU на месте.\n",
    "2. **Джойстик** — обучим ходьбу робота Go1: бегать вперёд, боком и поворачивать по «командам скорости».\n",
    "3. **Разбор полётов** — посмотрим видео походки и графики: как высоко робот поднимает лапы и точно ли держит скорость.\n",
    "4. **Стойка на передних лапах** — потренируем вторую политику и научимся дообучать её с сохранённого чекпоинта.\n",
    "5. **Бонус: двуногий робот** — если хватит времени и GPU.\n",
    "\n",
    "> Пока обучение идёт, посматривай на график награды: так выглядит «взросление» робота.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "Xqo7pyX-n72M"
   },
   "outputs": [],
   "source": [
    "#@title Устанавливаем главные библиотеки\n",
    "# Пояснение:\n",
    "#   mujoco     — симулятор физики: виртуальный мир, в котором живёт робот;\n",
    "#   mujoco_mjx — тот же MuJoCo, но на JAX: физика считается на GPU очень быстро;\n",
    "#   brax       — библиотека обучения роботов: в ней живёт алгоритм PPO,\n",
    "#                которым мы будем учить ходьбу.\n",
    "!pip install mujoco\n",
    "!pip install mujoco_mjx\n",
    "!pip install brax"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "fc179393"
   },
   "source": [
    "# Заметка: отдельное обновление flax/jax здесь не выполняется —\n",
    "# совместимость версий решает патч в ячейке «Временный патч…» перед запуском обучения.\n"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "IbZxYDxzoz5R"
   },
   "outputs": [],
   "source": [
    "# @title Проверяем, что GPU доступен и MuJoCo установился\n",
    "\n",
    "import distutils.util\n",
    "import os\n",
    "import subprocess\n",
    "\n",
    "# Проверяем, видно ли GPU: обучать нейросеть на CPU — очень-очень долго.\n",
    "if subprocess.run('nvidia-smi').returncode:\n",
    "  raise RuntimeError(\n",
    "      'Нет связи с GPU. Проверь, что выбран GPU-рантайм: '\n",
    "      'меню «Среда выполнения» → «Сменить среду выполнения» → GPU.'\n",
    "  )\n",
    "\n",
    "# Небольшая техническая настройка: добавляем ICD-конфиг, чтобы glvnd видел\n",
    "# EGL-драйвер NVIDIA (в Colab его ставят не через APT, поэтому конфига нет).\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",
    "    \"file_format_version\" : \"1.0.0\",\n",
    "    \"ICD\" : {\n",
    "        \"library_path\" : \"libEGL_nvidia.so.0\"\n",
    "    }\n",
    "}\n",
    "\"\"\")\n",
    "\n",
    "# Говорим MuJoCo рисовать картинку через GPU (EGL-бэкенд).\n",
    "print('Настраиваем GPU-рендеринг:')\n",
    "%env MUJOCO_GL=egl\n",
    "\n",
    "try:\n",
    "  print('Проверяем, что установка удалась:')\n",
    "  import mujoco\n",
    "\n",
    "  mujoco.MjModel.from_xml_string('<mujoco/>')\n",
    "except Exception as e:\n",
    "  raise e from RuntimeError(\n",
    "      'Что-то пошло не так при установке. Посмотри вывод команд выше.\\n'\n",
    "      'Если это Colab — проверь, что включён GPU-рантайм.'\n",
    "  )\n",
    "\n",
    "print('Установка прошла успешно.')\n",
    "\n",
    "# Просим XLA использовать Triton GEMM: на части GPU это ускоряет шаги примерно на 30%.\n",
    "xla_flags = os.environ.get('XLA_FLAGS', '')\n",
    "xla_flags += ' --xla_gpu_triton_gemm_any=True'\n",
    "os.environ['XLA_FLAGS'] = xla_flags"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "T5f4w3Kq2X14"
   },
   "outputs": [],
   "source": [
    "# @title Ставим и импортируем пакеты для графиков и видео\n",
    "import json\n",
    "import itertools\n",
    "import time\n",
    "from typing import Callable, List, NamedTuple, Optional, Union\n",
    "import numpy as np\n",
    "\n",
    "# Графики и видео.\n",
    "print(\"Устанавливаем mediapy:\")\n",
    "!command -v ffmpeg >/dev/null || (apt update && apt install -y ffmpeg)\n",
    "!pip install -q mediapy\n",
    "import mediapy as media\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Пусть numpy печатает числа компактно и по-человечески.\n",
    "np.set_printoptions(precision=3, suppress=True, linewidth=100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "ObF1UXrkb0Nd"
   },
   "outputs": [],
   "source": [
    "# @title Импортируем MuJoCo, MJX и Brax\n",
    "# Главные действующие лица:\n",
    "#   jax        — вычисления на GPU (на нём работает вся математика);\n",
    "#   mujoco/mjx — симулятор физики;\n",
    "#   brax       — обучение роботов (алгоритмы PPO/SAC);\n",
    "#   flax       — нейросети на JAX;\n",
    "#   orbax      — сохранение чекпоинтов (версий политики).\n",
    "from datetime import datetime\n",
    "import functools\n",
    "import os\n",
    "from typing import Any, Dict, Sequence, Tuple, Union\n",
    "from brax import base\n",
    "from brax import envs\n",
    "from brax import math\n",
    "from brax.base import Base, Motion, Transform\n",
    "from brax.base import State as PipelineState\n",
    "from brax.envs.base import Env, PipelineEnv, State\n",
    "from brax.io import html, mjcf, model\n",
    "from brax.mjx.base import State as MjxState\n",
    "from brax.training.agents.ppo import networks as ppo_networks\n",
    "from brax.training.agents.ppo import train as ppo\n",
    "from brax.training.agents.sac import networks as sac_networks\n",
    "from brax.training.agents.sac import train as sac\n",
    "from etils import epath\n",
    "from flax import struct\n",
    "from flax.training import orbax_utils\n",
    "from IPython.display import HTML, clear_output\n",
    "import jax\n",
    "from jax import numpy as jp\n",
    "from matplotlib import pyplot as plt\n",
    "import mediapy as media\n",
    "from ml_collections import config_dict\n",
    "import mujoco\n",
    "from mujoco import mjx\n",
    "import numpy as np\n",
    "from orbax import checkpoint as ocp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "UoTLSx4cFRdy"
   },
   "outputs": [],
   "source": [
    "#@title Ставим MuJoCo Playground\n",
    "# Playground — «игровая площадка» Google DeepMind:\n",
    "# коллекция готовых сред, настроек и обученных политик для роботов.\n",
    "!pip install playground"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "e620c19c"
   },
   "source": [
    "# Обновляем warp-lang — библиотеку быстрых GPU-расчётов,\n",
    "# которую используют внутренности Playground.\n",
    "!pip install --upgrade warp-lang"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "gYm2h7m8w3Nv"
   },
   "outputs": [],
   "source": [
    "#@title Импортируем The Playground\n",
    "# registry — «каталог»: по имени выдаёт среду, её настройки и функции.\n",
    "\n",
    "from mujoco_playground import wrapper\n",
    "from mujoco_playground import registry"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "LcibXbyKt4FI"
   },
   "source": [
    "# Ходьба (локомоция)\n",
    "\n",
    "В MuJoCo Playground собрано много сред для четвероногих и двуногих роботов. Полный список — в следующей ячейке.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ox0Gze9Ct5AM"
   },
   "outputs": [],
   "source": [
    "# Список всех сред для ходьбы, доступных в Playground.\n",
    "registry.locomotion.ALL_ENVS"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_R01tjWfI-i6"
   },
   "source": [
    "# Четвероногие\n",
    "\n",
    "Начнём сразу с четвероногих! В Playground есть среды для роботов Google Barkour и Boston Dynamics Spot, но лучше всего переносится на настоящего робота политика для **Unitree Go1** — того самого робота, которого ты мог видеть в лабораториях. Наша Панда устроена похожим образом, так что учимся на «двоюродном брате».\n",
    "\n",
    "Сначала обучим «джойстик»-политику: робот учится двигаться с той скоростью и в ту сторону, которую ему задают — как будто мы управляем им джойстиком.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "kPJeoQeEJBSA"
   },
   "outputs": [],
   "source": [
    "# Загружаем среду и её настройки по имени.\n",
    "env_name = 'Go1JoystickFlatTerrain'  # Go1: ровный пол, управление «джойстиком»\n",
    "env = registry.load(env_name)  # сама среда (симулятор)\n",
    "env_cfg = registry.get_default_config(env_name)  # настройки среды"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "6n9UT9N1wR5K"
   },
   "outputs": [],
   "source": [
    "# Смотрим настройки: команды, награды, длину эпизода и т.д.\n",
    "env_cfg"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Thm7nZueM4cz"
   },
   "source": [
    "## Джойстик\n",
    "\n",
    "Обучаем политику и смотрим, как она управляет роботом.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "B9T_UVZYLDdM"
   },
   "outputs": [],
   "source": [
    "# Достаём гиперпараметры обучения (алгоритм PPO) для нашей среды —\n",
    "# это «настройки тренера», которые мы можем менять.\n",
    "from mujoco_playground.config import locomotion_params\n",
    "ppo_params = locomotion_params.brax_ppo_config(env_name)\n",
    "ppo_params"
   ]
  },
  {
   "cell_type": "code",
   "source": [
    "# Наши настройки, чтобы уложиться по времени мастер-класса:\n",
    "ppo_params.num_timesteps = 80_000_000  # сколько шагов учится робот (по умолчанию 200_000_000)\n",
    "ppo_params.num_evals = 4  # сколько раз за обучение проверяем успехи (по умолчанию 10)\n",
    "ppo_params.episode_length = 1000  # длина одной попытки в шагах (по умолчанию 1000)"
   ],
   "metadata": {
    "id": "5BvZ-SKQXVRk"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Aefr2OS01D9g"
   },
   "source": [
    "При обучении используется **доменная рандомизация** (domain randomization) — приём, который помогает политике переноситься с симулятора на реального робота. Для части сред Playground уже есть готовые функции рандомизации: они лежат в registry и передаются прямо в алгоритмы обучения brax. Функция [доменной рандомизации](https://github.com/google-deepmind/mujoco_playground/blob/main/mujoco_playground/_src/locomotion/go1/randomize.py) случайно меняет трение, жёсткость приводов, центр масс корпуса, массы звеньев и другие параметры симуляции — чтобы робот учился ходить не в одном идеальном мире, а во множестве «неидеальных», как в реальности.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "UVA4Bn681DZT"
   },
   "outputs": [],
   "source": [
    "# Смотрим на функцию рандомизации для Go1 — её мы передадим в обучение.\n",
    "registry.get_domain_randomizer(env_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vBEEQyY6M5OC"
   },
   "source": [
    "### Обучение\n",
    "\n",
    "На RTX 4090 политика обучается за ~7 минут. В бесплатном Colab (GPU T4) дольше — успеем обсудить, что происходит с графиком награды, пока робот тренируется.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "XKFzyP7wM5OD"
   },
   "outputs": [],
   "source": [
    "# Готовим функцию обучения. progress() вызывается регулярно во время обучения\n",
    "# и обновляет график: как растёт средняя награда за эпизод.\n",
    "x_data, y_data, y_dataerr = [], [], []\n",
    "times = [datetime.now()]\n",
    "\n",
    "\n",
    "def progress(num_steps, metrics):\n",
    "  clear_output(wait=True)  # стираем прошлый график и рисуем обновлённый\n",
    "\n",
    "  times.append(datetime.now())\n",
    "  x_data.append(num_steps)\n",
    "  y_data.append(metrics[\"eval/episode_reward\"])\n",
    "  y_dataerr.append(metrics[\"eval/episode_reward_std\"])\n",
    "\n",
    "  plt.xlim([0, ppo_params[\"num_timesteps\"] * 1.25])\n",
    "  plt.xlabel(\"шаги обучения\")\n",
    "  plt.ylabel(\"награда за эпизод\")\n",
    "  plt.title(f\"награда = {y_data[-1]:.3f}\")\n",
    "  plt.errorbar(x_data, y_data, yerr=y_dataerr, color=\"blue\")\n",
    "\n",
    "  display(plt.gcf())\n",
    "\n",
    "# randomizer случайно «портит» условия (трение, массы…),\n",
    "# чтобы политика получилась устойчивой, как у настоящей Панды.\n",
    "randomizer = registry.get_domain_randomizer(env_name)\n",
    "# Собираем параметры обучения: гиперпараметры PPO + сеть + рандомизация + график.\n",
    "ppo_training_params = dict(ppo_params)\n",
    "network_factory = ppo_networks.make_ppo_networks\n",
    "if \"network_factory\" in ppo_params:\n",
    "  del ppo_training_params[\"network_factory\"]\n",
    "  network_factory = functools.partial(\n",
    "      ppo_networks.make_ppo_networks,\n",
    "      **ppo_params.network_factory\n",
    "  )\n",
    "\n",
    "train_fn = functools.partial(\n",
    "    ppo.train, **dict(ppo_training_params),\n",
    "    network_factory=network_factory,\n",
    "    randomization_fn=randomizer,\n",
    "    progress_fn=progress\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "FGrlulWbM5OD"
   },
   "outputs": [],
   "source": [
    "# Перед запуском — маленький патч совместимости версий:\n",
    "# новые версии JAX удалили функцию jax.device_put_replicated, которую использует\n",
    "# обучение. Подменяем её аналогом на современном API — чтобы всё заработало.\n",
    "# (Для мастер-класса: это «заплатка», суть — в запуске обучения ниже.)\n",
    "import jax\n",
    "import numpy as np\n",
    "\n",
    "def compatible_device_put_replicated(x, devices):\n",
    "    n_devices = len(devices)\n",
    "    # Создаём Mesh — «сетку» из доступных устройств (обычно это GPU).\n",
    "    mesh = jax.sharding.Mesh(np.array(devices), ('devices',))\n",
    "    # Схема распределения: делим массив по оси устройств.\n",
    "    sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec('devices'))\n",
    "\n",
    "    def replicate_leaf(leaf):\n",
    "        if isinstance(leaf, (int, float, bool)):\n",
    "            leaf = np.array(leaf)\n",
    "        # Копируем данные на каждое устройство.\n",
    "        stacked = np.stack([leaf] * n_devices)\n",
    "        # Размещаем через современный jax.device_put + NamedSharding.\n",
    "        return jax.device_put(stacked, sharding)\n",
    "\n",
    "    return jax.tree_util.tree_map(replicate_leaf, x)\n",
    "\n",
    "jax.device_put_replicated = compatible_device_put_replicated\n",
    "\n",
    "# Запускаем обучение! Это самый долгий шаг — следи за графиком награды.\n",
    "# make_inference_fn — «мозг» робота, params — его натренированные настройки.\n",
    "make_inference_fn, params, metrics = train_fn(\n",
    "    environment=env,\n",
    "    eval_env=registry.load(env_name, config=env_cfg),\n",
    "    wrap_env_fn=wrapper.wrap_for_brax_training,\n",
    ")\n",
    "print(f\"время компиляции (jit): {times[1] - times[0]}\")\n",
    "print(f\"время обучения: {times[-1] - times[1]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "AUxSNhq3UqmC"
   },
   "source": [
    "Смотрим, что получилось: запускаем обученную политику и рендерим видео!\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "RBM89g5A2Yoi"
   },
   "outputs": [],
   "source": [
    "# Включаем «толчки» (perturbation) в среде проверки — случайные толчки,\n",
    "# проверяющие, удержит ли робот равновесие.\n",
    "env_cfg = registry.get_default_config(env_name)\n",
    "env_cfg.pert_config.enable = True\n",
    "env_cfg.pert_config.velocity_kick = [3.0, 6.0]\n",
    "env_cfg.pert_config.kick_wait_times = [5.0, 15.0]\n",
    "env_cfg.command_config.a = [1.5, 0.8, 2*jp.pi]\n",
    "eval_env = registry.load(env_name, config=env_cfg)\n",
    "velocity_kick_range = [0.0, 0.0]  # Толчки «по скорости» отключены.\n",
    "kick_duration_range = [0.05, 0.2]  # Длительность толчка в секундах.\n",
    "\n",
    "# jax.jit компилирует функции под GPU — так они выполняются намного быстрее.\n",
    "jit_reset = jax.jit(eval_env.reset)\n",
    "jit_step = jax.jit(eval_env.step)\n",
    "jit_inference_fn = jax.jit(make_inference_fn(params, deterministic=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "C_1CY9xDoUKw"
   },
   "outputs": [],
   "source": [
    "#@title Запуск и рендер: видео походки\n",
    "from mujoco_playground._src.gait import draw_joystick_command\n",
    "\n",
    "# Команда роботу — можешь менять числа и запускать ячейку снова:\n",
    "x_vel = 0.0  #@param {type: \"number\"}\n",
    "y_vel = 1.0  #@param {type: \"number\"}\n",
    "yaw_vel = 0.0  #@param {type: \"number\"}\n",
    "# x_vel — вперёд/назад (м/с), y_vel — боком (м/с), yaw_vel — поворот (рад/с).\n",
    "\n",
    "\n",
    "def sample_pert(rng):\n",
    "  # Готовим случайный толчок: силу и длительность.\n",
    "  rng, key1, key2 = jax.random.split(rng, 3)\n",
    "  pert_mag = jax.random.uniform(\n",
    "      key1, minval=velocity_kick_range[0], maxval=velocity_kick_range[1]\n",
    "  )\n",
    "  duration_seconds = jax.random.uniform(\n",
    "      key2, minval=kick_duration_range[0], maxval=kick_duration_range[1]\n",
    "  )\n",
    "  duration_steps = jp.round(duration_seconds / eval_env.dt).astype(jp.int32)\n",
    "  state.info[\"pert_mag\"] = pert_mag\n",
    "  state.info[\"pert_duration\"] = duration_steps\n",
    "  state.info[\"pert_duration_seconds\"] = duration_seconds\n",
    "  return rng\n",
    "\n",
    "\n",
    "# Прогоняем один эпизод: робот действует по своей политике, а мы запоминаем\n",
    "# каждое его состояние — чтобы построить видео и графики.\n",
    "rng = jax.random.PRNGKey(0)\n",
    "rollout = []\n",
    "modify_scene_fns = []\n",
    "\n",
    "swing_peak = []\n",
    "rewards = []\n",
    "linvel = []\n",
    "angvel = []\n",
    "track = []\n",
    "foot_vel = []\n",
    "rews = []\n",
    "contact = []\n",
    "command = jp.array([x_vel, y_vel, yaw_vel])\n",
    "\n",
    "state = jit_reset(rng)\n",
    "if state.info[\"steps_since_last_pert\"] < state.info[\"steps_until_next_pert\"]:\n",
    "  rng = sample_pert(rng)\n",
    "state.info[\"command\"] = command\n",
    "for i in range(env_cfg.episode_length):\n",
    "  if state.info[\"steps_since_last_pert\"] < state.info[\"steps_until_next_pert\"]:\n",
    "    rng = sample_pert(rng)\n",
    "  act_rng, rng = jax.random.split(rng)\n",
    "  ctrl, _ = jit_inference_fn(state.obs, act_rng)\n",
    "  state = jit_step(state, ctrl)\n",
    "  state.info[\"command\"] = command\n",
    "  rews.append(\n",
    "      {k: v for k, v in state.metrics.items() if k.startswith(\"reward/\")}\n",
    "  )\n",
    "  rollout.append(state)\n",
    "  swing_peak.append(state.info[\"swing_peak\"])\n",
    "  rewards.append(\n",
    "      {k[7:]: v for k, v in state.metrics.items() if k.startswith(\"reward/\")}\n",
    "  )\n",
    "  linvel.append(env.get_global_linvel(state.data))\n",
    "  angvel.append(env.get_gyro(state.data))\n",
    "  track.append(\n",
    "      env._reward_tracking_lin_vel(\n",
    "          state.info[\"command\"], env.get_local_linvel(state.data)\n",
    "      )\n",
    "  )\n",
    "\n",
    "  feet_vel = state.data.sensordata[env._foot_linvel_sensor_adr]\n",
    "  vel_xy = feet_vel[..., :2]\n",
    "  vel_norm = jp.sqrt(jp.linalg.norm(vel_xy, axis=-1))\n",
    "  foot_vel.append(vel_norm)\n",
    "\n",
    "  contact.append(state.info[\"last_contact\"])\n",
    "\n",
    "  # Рисуем над роботом стрелку-команду: в видео видно, куда он «должен» бежать.\n",
    "  xyz = np.array(state.data.xpos[env._torso_body_id])\n",
    "  xyz += np.array([0, 0, 0.2])\n",
    "  x_axis = state.data.xmat[env._torso_body_id, 0]\n",
    "  yaw = -np.arctan2(x_axis[1], x_axis[0])\n",
    "  modify_scene_fns.append(\n",
    "      functools.partial(\n",
    "          draw_joystick_command,\n",
    "          cmd=state.info[\"command\"],\n",
    "          xyz=xyz,\n",
    "          theta=yaw,\n",
    "          scl=abs(state.info[\"command\"][0])\n",
    "          / env_cfg.command_config.a[0],\n",
    "      )\n",
    "  )\n",
    "\n",
    "\n",
    "# В видео берём каждый второй шаг — файл меньше, а смотрится так же.\n",
    "render_every = 2\n",
    "fps = 1.0 / eval_env.dt / render_every\n",
    "traj = rollout[::render_every]\n",
    "mod_fns = modify_scene_fns[::render_every]\n",
    "\n",
    "scene_option = mujoco.MjvOption()\n",
    "scene_option.geomgroup[2] = True\n",
    "scene_option.geomgroup[3] = False\n",
    "scene_option.flags[mujoco.mjtVisFlag.mjVIS_CONTACTPOINT] = True\n",
    "scene_option.flags[mujoco.mjtVisFlag.mjVIS_TRANSPARENT] = False\n",
    "scene_option.flags[mujoco.mjtVisFlag.mjVIS_PERTFORCE] = True\n",
    "\n",
    "frames = eval_env.render(\n",
    "    traj,\n",
    "    camera=\"track\",\n",
    "    scene_option=scene_option,\n",
    "    width=640,\n",
    "    height=480,\n",
    "    modify_scene_fns=mod_fns,\n",
    ")\n",
    "media.show_video(frames, fps=fps, loop=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1QHdoJ2r30En"
   },
   "source": [
    "Посмотрим на графики: как поднимаются лапы (какая походка получилась) и держит ли робот заданную скорость. В лабораториях по таким графикам отлаживают политику — как спортсмен смотрит свой пульс и время на кругах!\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "gyyynm3ozEet"
   },
   "outputs": [],
   "source": [
    "#@title Графики: высота лап и точность следования команде\n",
    "\n",
    "# Лапы: FR — передняя правая, FL — передняя левая, RR — задняя правая, RL — задняя левая.\n",
    "swing_peak = jp.array(swing_peak)\n",
    "names = [\"FR\", \"FL\", \"RR\", \"RL\"]\n",
    "colors = [\"r\", \"g\", \"b\", \"y\"]\n",
    "fig, axs = plt.subplots(2, 2)\n",
    "for i, ax in enumerate(axs.flat):\n",
    "  ax.plot(swing_peak[:, i], color=colors[i])\n",
    "  ax.set_ylim([0, env_cfg.reward_config.max_foot_height * 1.25])\n",
    "  ax.axhline(env_cfg.reward_config.max_foot_height, color=\"k\", linestyle=\"--\")\n",
    "  ax.set_title(names[i])\n",
    "  ax.set_xlabel(\"время\")\n",
    "  ax.set_ylabel(\"высота\")\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "linvel_x = jp.array(linvel)[:, 0]\n",
    "linvel_y = jp.array(linvel)[:, 1]\n",
    "angvel_yaw = jp.array(angvel)[:, 2]\n",
    "\n",
    "# Сглаживаем графики скользящим средним: так видны тенденции, а не шум.\n",
    "linvel_x = jp.convolve(linvel_x, jp.ones(10) / 10, mode=\"same\")\n",
    "linvel_y = jp.convolve(linvel_y, jp.ones(10) / 10, mode=\"same\")\n",
    "angvel_yaw = jp.convolve(angvel_yaw, jp.ones(10) / 10, mode=\"same\")\n",
    "\n",
    "# Красная пунктирная линия — скорость, которую просили у робота;\n",
    "# синяя — скорость, которую он выдал на самом деле.\n",
    "fig, axes = plt.subplots(3, 1, figsize=(10, 10))\n",
    "axes[0].plot(linvel_x)\n",
    "axes[1].plot(linvel_y)\n",
    "axes[2].plot(angvel_yaw)\n",
    "\n",
    "axes[0].set_ylim(\n",
    "    -env_cfg.command_config.a[0], env_cfg.command_config.a[0]\n",
    ")\n",
    "axes[1].set_ylim(\n",
    "    -env_cfg.command_config.a[1], env_cfg.command_config.a[1]\n",
    ")\n",
    "axes[2].set_ylim(\n",
    "    -env_cfg.command_config.a[2], env_cfg.command_config.a[2]\n",
    ")\n",
    "\n",
    "for i, ax in enumerate(axes):\n",
    "  ax.axhline(state.info[\"command\"][i], color=\"red\", linestyle=\"--\")\n",
    "\n",
    "labels = [\"dx\", \"dy\", \"dyaw\"]  # вперёд/назад, боком, поворот\n",
    "for i, ax in enumerate(axes):\n",
    "  ax.set_ylabel(labels[i])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "t1QAHuYBQBbl"
   },
   "source": [
    "А теперь — плавный разгон: каждые 200 шагов команда скорости увеличивается на 0,25 м/с. Смотри видео, как робот набирает ход.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "Q0EuQiVlzh5u"
   },
   "outputs": [],
   "source": [
    "#@title Плавно увеличиваем скорость вперёд\n",
    "\n",
    "rng = jax.random.PRNGKey(0)\n",
    "rollout = []\n",
    "modify_scene_fns = []\n",
    "swing_peak = []\n",
    "linvel = []\n",
    "angvel = []\n",
    "\n",
    "x = -0.25\n",
    "command = jp.array([x, 0, 0])\n",
    "state = jit_reset(rng)\n",
    "for i in range(1_400):\n",
    "  # Каждые 200 шагов прибавляем 0,25 м/с к скорости вперёд.\n",
    "  if i % 200 == 0:\n",
    "    x += 0.25\n",
    "    print(f\"Задаю скорость вперёд: {x} м/с\")\n",
    "    command = jp.array([x, 0, 0])\n",
    "  state.info[\"command\"] = command\n",
    "  if state.info[\"steps_since_last_pert\"] < state.info[\"steps_until_next_pert\"]:\n",
    "    rng = sample_pert(rng)\n",
    "  act_rng, rng = jax.random.split(rng)\n",
    "  ctrl, _ = jit_inference_fn(state.obs, act_rng)\n",
    "  state = jit_step(state, ctrl)\n",
    "  rollout.append(state)\n",
    "  swing_peak.append(state.info[\"swing_peak\"])\n",
    "  linvel.append(env.get_global_linvel(state.data))\n",
    "  angvel.append(env.get_gyro(state.data))\n",
    "  xyz = np.array(state.data.xpos[env._torso_body_id])\n",
    "  xyz += np.array([0, 0, 0.2])\n",
    "  x_axis = state.data.xmat[env._torso_body_id, 0]\n",
    "  yaw = -np.arctan2(x_axis[1], x_axis[0])\n",
    "  modify_scene_fns.append(\n",
    "      functools.partial(\n",
    "          draw_joystick_command,\n",
    "          cmd=command,\n",
    "          xyz=xyz,\n",
    "          theta=yaw,\n",
    "          scl=abs(command[0]) / env_cfg.command_config.a[0],\n",
    "      )\n",
    "  )\n",
    "\n",
    "\n",
    "# Графики высоты лап — как в предыдущей проверке.\n",
    "swing_peak = jp.array(swing_peak)\n",
    "names = [\"FR\", \"FL\", \"RR\", \"RL\"]\n",
    "colors = [\"r\", \"g\", \"b\", \"y\"]\n",
    "fig, axs = plt.subplots(2, 2)\n",
    "for i, ax in enumerate(axs.flat):\n",
    "  ax.plot(swing_peak[:, i], color=colors[i])\n",
    "  ax.set_ylim([0, env_cfg.reward_config.max_foot_height * 1.25])\n",
    "  ax.axhline(env_cfg.reward_config.max_foot_height, color=\"k\", linestyle=\"--\")\n",
    "  ax.set_title(names[i])\n",
    "  ax.set_xlabel(\"время\")\n",
    "  ax.set_ylabel(\"высота\")\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "linvel_x = jp.array(linvel)[:, 0]\n",
    "linvel_y = jp.array(linvel)[:, 1]\n",
    "angvel_yaw = jp.array(angvel)[:, 2]\n",
    "\n",
    "# Сверяем реальную скорость с командой.\n",
    "linvel_x = jp.convolve(linvel_x, jp.ones(10) / 10, mode=\"same\")\n",
    "linvel_y = jp.convolve(linvel_y, jp.ones(10) / 10, mode=\"same\")\n",
    "angvel_yaw = jp.convolve(angvel_yaw, jp.ones(10) / 10, mode=\"same\")\n",
    "\n",
    "fig, axes = plt.subplots(3, 1, figsize=(10, 10))\n",
    "axes[0].plot(linvel_x)\n",
    "axes[1].plot(linvel_y)\n",
    "axes[2].plot(angvel_yaw)\n",
    "\n",
    "axes[0].set_ylim(\n",
    "    -env_cfg.command_config.a[0], env_cfg.command_config.a[0]\n",
    ")\n",
    "axes[1].set_ylim(\n",
    "    -env_cfg.command_config.a[1], env_cfg.command_config.a[1]\n",
    ")\n",
    "axes[2].set_ylim(\n",
    "    -env_cfg.command_config.a[2], env_cfg.command_config.a[2]\n",
    ")\n",
    "\n",
    "for i, ax in enumerate(axes):\n",
    "  ax.axhline(state.info[\"command\"][i], color=\"red\", linestyle=\"--\")\n",
    "\n",
    "labels = [\"dx\", \"dy\", \"dyaw\"]\n",
    "for i, ax in enumerate(axes):\n",
    "  ax.set_ylabel(labels[i])\n",
    "\n",
    "\n",
    "render_every = 2\n",
    "fps = 1.0 / eval_env.dt / render_every\n",
    "print(f\"fps: {fps}\")\n",
    "\n",
    "traj = rollout[::render_every]\n",
    "mod_fns = modify_scene_fns[::render_every]\n",
    "assert len(traj) == len(mod_fns)\n",
    "\n",
    "scene_option = mujoco.MjvOption()\n",
    "scene_option.geomgroup[2] = True\n",
    "scene_option.geomgroup[3] = False\n",
    "scene_option.flags[mujoco.mjtVisFlag.mjVIS_CONTACTPOINT] = True\n",
    "scene_option.flags[mujoco.mjtVisFlag.mjVIS_PERTFORCE] = True\n",
    "\n",
    "frames = eval_env.render(\n",
    "    traj,\n",
    "    camera=\"track\",\n",
    "    height=480,\n",
    "    width=640,\n",
    "    modify_scene_fns=mod_fns,\n",
    "    scene_option=scene_option,\n",
    ")\n",
    "media.show_video(frames, fps=fps, loop=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0RHZvXgmzrEJ"
   },
   "source": [
    "## Стойка на передних лапах (handstand)\n",
    "\n",
    "Для Unitree Go1 есть и другие политики: подъём после падения, стойка на передних лапах, стойка на всех лапах и т.д. На примере стойки покажем, как **дообучать** политику с сохранённого чекпоинта — так можно быстро переучивать робота, меняя настройки среды между запусками.\n",
    "\n",
    "Для стойки сначала обучимся с обычными настройками, а затем добавим штраф за энергию — политика станет плавнее и будет лучше переноситься на реального робота.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "RYriZOAxzEk_"
   },
   "outputs": [],
   "source": [
    "# Схема как раньше, но среда другая: учим стойку на передних лапах.\n",
    "from mujoco_playground.config import locomotion_params\n",
    "\n",
    "env_name = 'Go1Handstand'\n",
    "env = registry.load(env_name)\n",
    "env_cfg = registry.get_default_config(env_name)\n",
    "ppo_params = locomotion_params.brax_ppo_config(env_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3nB5ugbdS5kk"
   },
   "source": [
    "Создаём каталог для чекпоинтов — «сохранённых версий» политики, и учимся с сохранением.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "EyEDpHisS7eO"
   },
   "outputs": [],
   "source": [
    "# Каталог для чекпоинтов — «сохранённых версий» политики.\n",
    "ckpt_path = epath.Path(\"checkpoints\").resolve() / env_name\n",
    "ckpt_path.mkdir(parents=True, exist_ok=True)\n",
    "print(f\"{ckpt_path}\")\n",
    "\n",
    "# Рядом сохраняем конфиг среды — чтобы помнить, с какими настройками учились.\n",
    "with open(ckpt_path / \"config.json\", \"w\") as fp:\n",
    "  json.dump(env_cfg.to_dict(), fp, indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "lCRUYofXSNGT"
   },
   "outputs": [],
   "source": [
    "#@title Определяем функцию обучения (с сохранением чекпоинтов)\n",
    "x_data, y_data, y_dataerr = [], [], []\n",
    "times = [datetime.now()]\n",
    "\n",
    "\n",
    "def policy_params_fn(current_step, make_policy, params):\n",
    "  del make_policy  # Не используется.\n",
    "  # Кладём текущие параметры политики в каталог чекпоинтов.\n",
    "  orbax_checkpointer = ocp.PyTreeCheckpointer()\n",
    "  save_args = orbax_utils.save_args_from_target(params)\n",
    "  path = ckpt_path / f\"{current_step}\"\n",
    "  orbax_checkpointer.save(path, params, force=True, save_args=save_args)\n",
    "\n",
    "\n",
    "def progress(num_steps, metrics):\n",
    "  clear_output(wait=True)\n",
    "\n",
    "  times.append(datetime.now())\n",
    "  x_data.append(num_steps)\n",
    "  y_data.append(metrics[\"eval/episode_reward\"])\n",
    "  y_dataerr.append(metrics[\"eval/episode_reward_std\"])\n",
    "\n",
    "  plt.xlim([0, ppo_params[\"num_timesteps\"] * 1.25])\n",
    "  plt.xlabel(\"шаги обучения\")\n",
    "  plt.ylabel(\"награда за эпизод\")\n",
    "  plt.title(f\"награда = {y_data[-1]:.3f}\")\n",
    "  plt.errorbar(x_data, y_data, yerr=y_dataerr, color=\"blue\")\n",
    "\n",
    "  display(plt.gcf())\n",
    "\n",
    "randomizer = registry.get_domain_randomizer(env_name)\n",
    "ppo_training_params = dict(ppo_params)\n",
    "network_factory = ppo_networks.make_ppo_networks\n",
    "if \"network_factory\" in ppo_params:\n",
    "  del ppo_training_params[\"network_factory\"]\n",
    "  network_factory = functools.partial(\n",
    "      ppo_networks.make_ppo_networks,\n",
    "      **ppo_params.network_factory\n",
    "  )\n",
    "\n",
    "train_fn = functools.partial(\n",
    "    ppo.train, **dict(ppo_training_params),\n",
    "    network_factory=network_factory,\n",
    "    randomization_fn=randomizer,\n",
    "    progress_fn=progress,\n",
    "    policy_params_fn=policy_params_fn,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "A1oK80x1anPp"
   },
   "source": [
    "Первое обучение стойки занимает ~8 минут на RTX 4090.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "MY6P3abhSNGU"
   },
   "outputs": [],
   "source": [
    "# Запускаем обучение стойки. По ходу сохраняются чекпоинты.\n",
    "make_inference_fn, params, metrics = train_fn(\n",
    "    environment=registry.load(env_name, config=env_cfg),\n",
    "    eval_env=registry.load(env_name, config=env_cfg),\n",
    "    wrap_env_fn=wrapper.wrap_for_brax_training,\n",
    ")\n",
    "print(f\"время компиляции (jit): {times[1] - times[0]}\")\n",
    "print(f\"время обучения: {times[-1] - times[1]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4s6PkZ4GWV4Z"
   },
   "source": [
    "Смотрим текущую политику — до дообучения.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "WiWOtc_6WbcX"
   },
   "outputs": [],
   "source": [
    "#@title Запуск и рендер (стойка)\n",
    "inference_fn = make_inference_fn(params, deterministic=True)\n",
    "jit_inference_fn = jax.jit(inference_fn)\n",
    "\n",
    "eval_env = registry.load(env_name, config=env_cfg)\n",
    "jit_reset = jax.jit(eval_env.reset)\n",
    "jit_step = jax.jit(eval_env.step)\n",
    "\n",
    "rng = jax.random.PRNGKey(12345)\n",
    "rollout = []\n",
    "rewards = []\n",
    "torso_height = []\n",
    "actions = []\n",
    "torques = []\n",
    "power = []\n",
    "qfrc_constraint = []\n",
    "qvels = []\n",
    "power1 = []\n",
    "power2 = []\n",
    "# Прогоняем 10 эпизодов и собираем «цену» движения: силы моторов,\n",
    "# мощность, ускорения — чтобы понять, насколько политика плавная.\n",
    "for _ in range(10):\n",
    "  rng, reset_rng = jax.random.split(rng)\n",
    "  state = jit_reset(reset_rng)\n",
    "  for i in range(env_cfg.episode_length // 2):\n",
    "    act_rng, rng = jax.random.split(rng)\n",
    "    ctrl, _ = jit_inference_fn(state.obs, act_rng)\n",
    "    actions.append(ctrl)\n",
    "    state = jit_step(state, ctrl)\n",
    "    rollout.append(state)\n",
    "    rewards.append(\n",
    "        {k[7:]: v for k, v in state.metrics.items() if k.startswith(\"reward/\")}\n",
    "    )\n",
    "    torso_height.append(state.data.qpos[2])\n",
    "    torques.append(state.data.actuator_force)\n",
    "    qvel = state.data.qvel[6:]\n",
    "    power.append(jp.sum(jp.abs(qvel * state.data.actuator_force)))\n",
    "    qfrc_constraint.append(jp.linalg.norm(state.data.qfrc_constraint[6:]))\n",
    "    qvels.append(jp.max(jp.abs(qvel)))\n",
    "    frc = state.data.actuator_force\n",
    "    qvel = state.data.qvel[6:]\n",
    "    power1.append(jp.sum(frc * qvel))\n",
    "    power2.append(jp.sum(jp.abs(frc * qvel)))\n",
    "\n",
    "\n",
    "render_every = 2\n",
    "fps = 1.0 / eval_env.dt / render_every\n",
    "traj = rollout[::render_every]\n",
    "\n",
    "scene_option = mujoco.MjvOption()\n",
    "scene_option.geomgroup[2] = True\n",
    "scene_option.geomgroup[3] = False\n",
    "scene_option.flags[mujoco.mjtVisFlag.mjVIS_CONTACTPOINT] = True\n",
    "scene_option.flags[mujoco.mjtVisFlag.mjVIS_CONTACTFORCE] = False\n",
    "scene_option.flags[mujoco.mjtVisFlag.mjVIS_TRANSPARENT] = False\n",
    "\n",
    "frames = eval_env.render(\n",
    "    traj, camera=\"side\", scene_option=scene_option, height=480, width=640\n",
    ")\n",
    "media.show_video(frames, fps=fps, loop=False)\n",
    "\n",
    "power = jp.array(power1)\n",
    "print(f\"Максимальная мощность: {jp.max(power)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "v5p0Z3PPSRik"
   },
   "source": [
    "Заметил, что политика дёргается? На реального робота такую переносить нельзя — и максимальная мощность тоже слишком высокая.\n",
    "\n",
    "Для переноса в реальность стойку учили с «учебным планом»: настройки `energy_termination_threshold`, `energy` и `dof_acc` штрафуют высокие моменты и большую мощность. Давай дообучим политику с пониженным `energy_termination_threshold` и ненулевыми штрафами `energy` и `dof_acc` — движения должны стать плавнее.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hrjoVL-_WN-r"
   },
   "source": [
    "### Дообучаем с последнего чекпоинта\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "jTxAySRSSu96"
   },
   "outputs": [],
   "source": [
    "env_cfg = registry.get_default_config(env_name)\n",
    "env_cfg.energy_termination_threshold = 400  # понижаем порог остановки по энергии\n",
    "env_cfg.reward_config.energy = -0.003  # штраф (отрицательная награда) за энергию\n",
    "env_cfg.reward_config.dof_acc = -2.5e-7  # штраф за ускорение приводов\n",
    "\n",
    "# Ищем самый свежий чекпоинт — дообучение продолжится с него.\n",
    "FINETUNE_PATH = epath.Path(ckpt_path)\n",
    "latest_ckpts = list(FINETUNE_PATH.glob(\"*\"))\n",
    "latest_ckpts = [ckpt for ckpt in latest_ckpts if ckpt.is_dir()]\n",
    "latest_ckpts.sort(key=lambda x: int(x.name))\n",
    "latest_ckpt = latest_ckpts[-1]\n",
    "restore_checkpoint_path = latest_ckpt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_M5IqOR6z4bV"
   },
   "outputs": [],
   "source": [
    "x_data, y_data, y_dataerr = [], [], []\n",
    "# Дообучаем с сохранённого чекпоинта и фиксируем seed=1,\n",
    "# чтобы результат повторялся при следующих запусках.\n",
    "times = [datetime.now()]\n",
    "\n",
    "make_inference_fn, params, metrics = train_fn(\n",
    "    environment=registry.load(env_name, config=env_cfg),\n",
    "    eval_env=registry.load(env_name, config=env_cfg),\n",
    "    wrap_env_fn=wrapper.wrap_for_brax_training,\n",
    "    restore_checkpoint_path=restore_checkpoint_path,  # продолжаем с чекпоинта!\n",
    "    seed=1,\n",
    ")\n",
    "print(f\"время компиляции (jit): {times[1] - times[0]}\")\n",
    "print(f\"время обучения: {times[-1] - times[1]}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "tzG8eY2lz4dk"
   },
   "outputs": [],
   "source": [
    "#@title Запуск и рендер (дообученная стойка)\n",
    "inference_fn = make_inference_fn(params, deterministic=True)\n",
    "jit_inference_fn = jax.jit(inference_fn)\n",
    "\n",
    "eval_env = registry.load(env_name, config=env_cfg)\n",
    "jit_reset = jax.jit(eval_env.reset)\n",
    "jit_step = jax.jit(eval_env.step)\n",
    "\n",
    "rng = jax.random.PRNGKey(12345)\n",
    "rollout = []\n",
    "rewards = []\n",
    "torso_height = []\n",
    "actions = []\n",
    "torques = []\n",
    "power = []\n",
    "qfrc_constraint = []\n",
    "qvels = []\n",
    "power1 = []\n",
    "power2 = []\n",
    "# Прогоняем 10 эпизодов: сравним плавность до и после дообучения.\n",
    "for _ in range(10):\n",
    "  rng, reset_rng = jax.random.split(rng)\n",
    "  state = jit_reset(reset_rng)\n",
    "  for i in range(env_cfg.episode_length // 2):\n",
    "    act_rng, rng = jax.random.split(rng)\n",
    "    ctrl, _ = jit_inference_fn(state.obs, act_rng)\n",
    "    actions.append(ctrl)\n",
    "    state = jit_step(state, ctrl)\n",
    "    rollout.append(state)\n",
    "    rewards.append(\n",
    "        {k[7:]: v for k, v in state.metrics.items() if k.startswith(\"reward/\")}\n",
    "    )\n",
    "    torso_height.append(state.data.qpos[2])\n",
    "    torques.append(state.data.actuator_force)\n",
    "    qvel = state.data.qvel[6:]\n",
    "    power.append(jp.sum(jp.abs(qvel * state.data.actuator_force)))\n",
    "    qfrc_constraint.append(jp.linalg.norm(state.data.qfrc_constraint[6:]))\n",
    "    qvels.append(jp.max(jp.abs(qvel)))\n",
    "    frc = state.data.actuator_force\n",
    "    qvel = state.data.qvel[6:]\n",
    "    power1.append(jp.sum(frc * qvel))\n",
    "    power2.append(jp.sum(jp.abs(frc * qvel)))\n",
    "\n",
    "\n",
    "render_every = 2\n",
    "fps = 1.0 / eval_env.dt / render_every\n",
    "traj = rollout[::render_every]\n",
    "\n",
    "scene_option = mujoco.MjvOption()\n",
    "scene_option.geomgroup[2] = True\n",
    "scene_option.geomgroup[3] = False\n",
    "scene_option.flags[mujoco.mjtVisFlag.mjVIS_CONTACTPOINT] = True\n",
    "scene_option.flags[mujoco.mjtVisFlag.mjVIS_CONTACTFORCE] = False\n",
    "scene_option.flags[mujoco.mjtVisFlag.mjVIS_TRANSPARENT] = False\n",
    "\n",
    "frames = eval_env.render(\n",
    "    traj, camera=\"side\", scene_option=scene_option, height=480, width=640\n",
    ")\n",
    "media.show_video(frames, fps=fps, loop=False)\n",
    "\n",
    "power = jp.array(power1)\n",
    "print(f\"Максимальная мощность: {jp.max(power)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yCyibqGMiAca"
   },
   "source": [
    "Теперь политика должна двигаться плавнее и потреблять меньше мощности! Попробуй дообучить её ещё с другими штрафами — и сравнить, что получится.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "26o77FfWXvVp"
   },
   "source": [
    "# Двуногие\n",
    "\n",
    "В Playground есть и двуногие роботы: Berkeley Humanoid, Unitree G1/H1. Покажем «джойстик»-политику для Berkeley Humanoid. Первое обучение занимает ~17 минут на RTX 4090 — это бонусная часть, если осталось время и GPU.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ESNd18FUanPt"
   },
   "outputs": [],
   "source": [
    "# Двуногий робот Berkeley Humanoid. Схема — точно как с Go1!\n",
    "env_name = 'BerkeleyHumanoidJoystickFlatTerrain'\n",
    "env = registry.load(env_name)\n",
    "env_cfg = registry.get_default_config(env_name)\n",
    "ppo_params = locomotion_params.brax_ppo_config(env_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "nibLoRu8anPt"
   },
   "outputs": [],
   "source": [
    "# Собираем функцию обучения — всё как раньше.\n",
    "x_data, y_data, y_dataerr = [], [], []\n",
    "times = [datetime.now()]\n",
    "\n",
    "randomizer = registry.get_domain_randomizer(env_name)\n",
    "ppo_training_params = dict(ppo_params)\n",
    "network_factory = ppo_networks.make_ppo_networks\n",
    "if \"network_factory\" in ppo_params:\n",
    "  del ppo_training_params[\"network_factory\"]\n",
    "  network_factory = functools.partial(\n",
    "      ppo_networks.make_ppo_networks,\n",
    "      **ppo_params.network_factory\n",
    "  )\n",
    "\n",
    "train_fn = functools.partial(\n",
    "    ppo.train, **dict(ppo_training_params),\n",
    "    network_factory=network_factory,\n",
    "    randomization_fn=randomizer,\n",
    "    progress_fn=progress\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "16dqomv0anPt"
   },
   "outputs": [],
   "source": [
    "# Запускаем обучение двуногого робота.\n",
    "make_inference_fn, params, metrics = train_fn(\n",
    "    environment=env,\n",
    "    eval_env=registry.load(env_name, config=env_cfg),\n",
    "    wrap_env_fn=wrapper.wrap_for_brax_training,\n",
    ")\n",
    "print(f\"время компиляции (jit): {times[1] - times[0]}\")\n",
    "print(f\"время обучения: {times[-1] - times[1]}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "sBHDF-JFanPt",
    "cellView": "form"
   },
   "outputs": [],
   "source": [
    "#@title Запуск и рендер (двуногий робот)\n",
    "from mujoco_playground._src.gait import draw_joystick_command\n",
    "\n",
    "env = registry.load(env_name)\n",
    "eval_env = registry.load(env_name)\n",
    "jit_reset = jax.jit(eval_env.reset)\n",
    "jit_step = jax.jit(eval_env.step)\n",
    "jit_inference_fn = jax.jit(make_inference_fn(params, deterministic=True))\n",
    "\n",
    "rng = jax.random.PRNGKey(1)\n",
    "\n",
    "rollout = []\n",
    "modify_scene_fns = []\n",
    "\n",
    "# Команда двуногому: вперёд (x), боком (y), поворот (yaw).\n",
    "x_vel = 1.0  #@param {type: \"number\"}\n",
    "y_vel = 0.0  #@param {type: \"number\"}\n",
    "yaw_vel = 0.0  #@param {type: \"number\"}\n",
    "command = jp.array([x_vel, y_vel, yaw_vel])\n",
    "\n",
    "# phase — «фаза шага»: ритм, в котором робот переставляет ноги.\n",
    "phase_dt = 2 * jp.pi * eval_env.dt * 1.5\n",
    "phase = jp.array([0, jp.pi])\n",
    "\n",
    "for j in range(1):\n",
    "  print(f\"эпизод {j}\")\n",
    "  state = jit_reset(rng)\n",
    "  state.info[\"phase_dt\"] = phase_dt\n",
    "  state.info[\"phase\"] = phase\n",
    "  for i in range(env_cfg.episode_length):\n",
    "    act_rng, rng = jax.random.split(rng)\n",
    "    ctrl, _ = jit_inference_fn(state.obs, act_rng)\n",
    "    state = jit_step(state, ctrl)\n",
    "    if state.done:\n",
    "      break\n",
    "    state.info[\"command\"] = command\n",
    "    rollout.append(state)\n",
    "\n",
    "    xyz = np.array(state.data.xpos[eval_env.mj_model.body(\"torso\").id])\n",
    "    xyz += np.array([0, 0.0, 0])\n",
    "    x_axis = state.data.xmat[eval_env._torso_body_id, 0]\n",
    "    yaw = -np.arctan2(x_axis[1], x_axis[0])\n",
    "    modify_scene_fns.append(\n",
    "        functools.partial(\n",
    "            draw_joystick_command,\n",
    "            cmd=state.info[\"command\"],\n",
    "            xyz=xyz,\n",
    "            theta=yaw,\n",
    "            scl=np.linalg.norm(state.info[\"command\"]),\n",
    "        )\n",
    "    )\n",
    "\n",
    "render_every = 1\n",
    "fps = 1.0 / eval_env.dt / render_every\n",
    "print(f\"fps: {fps}\")\n",
    "traj = rollout[::render_every]\n",
    "mod_fns = modify_scene_fns[::render_every]\n",
    "\n",
    "scene_option = mujoco.MjvOption()\n",
    "scene_option.geomgroup[2] = True\n",
    "scene_option.geomgroup[3] = False\n",
    "scene_option.flags[mujoco.mjtVisFlag.mjVIS_CONTACTPOINT] = True\n",
    "scene_option.flags[mujoco.mjtVisFlag.mjVIS_TRANSPARENT] = False\n",
    "scene_option.flags[mujoco.mjtVisFlag.mjVIS_PERTFORCE] = False\n",
    "\n",
    "frames = eval_env.render(\n",
    "    traj,\n",
    "    camera=\"track\",\n",
    "    scene_option=scene_option,\n",
    "    width=640*2,\n",
    "    height=480,\n",
    "    modify_scene_fns=mod_fns,\n",
    ")\n",
    "media.show_video(frames, fps=fps, loop=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CBtrAqns35sI"
   },
   "source": [
    "🙌 До скорой встречи!\n"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "private_outputs": true,
   "provenance": [],
   "toc_visible": true,
   "machine_shape": "hm",
   "gpuType": "T4",
   "name": "Плейбук — учим робота ходить"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  },
  "accelerator": "GPU"
 },
 "nbformat": 4,
 "nbformat_minor": 0
}