simulate PD controller for A1 based on data

examples/a1_pd_render.ipynb

@@ -0,0 +1,120 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import mujoco\n",
+    "import mediapy as media\n",
+    "import numpy as np \n",
+    "from robot_descriptions.a1_mj_description import MJCF_PATH\n",
+    "from scipy import interpolate\n",
+    "from utils import load_model_with_scene"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "model = load_model_with_scene(MJCF_PATH)\n",
+    "renderer = mujoco.Renderer(model, height=480, width=640)#, width=800, height=600)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "data_file = 'run_twist_active.csv'\n",
+    "\n",
+    "data_array = np.genfromtxt(data_file, delimiter=',', skip_header=100, skip_footer=100)\n",
+    "timespan = data_array[:, 0]- data_array[0, 0]\n",
+    "# sampling = np.mean(np.diff(timespan))\n",
+    "quaternion = data_array[:, 1:5]\n",
+    "joint_angles = data_array[:, 11:23] \n",
+    "# linear_velocity = data_array[:, 5:8]\n",
+    "# joint_velocity = data_array[:, 23:35] \n",
+    "# joint_torques = data_array[:, 35:47]\n",
+    "# contacts = data_array[:, 47:]\n",
+    "final_time = timespan[-1]\n",
+    "\n",
+    "\n",
+    "joint_angles_t = interpolate.interp1d(timespan, joint_angles, axis = 0)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "data = mujoco.MjData(model)\n",
+    "mujoco.mj_resetData(model, data)  # Reset state and time.\n",
+    "mujoco.mj_resetData(model, data)  # Reset state and time.\n",
+    "mujoco.mj_step(model, data)\n",
+    "model.opt.timestep = 1/1000\n",
+    "\n",
+    "scene_option = mujoco.MjvOption()\n",
+    "scene_option.flags[mujoco.mjtVisFlag.mjVIS_CONTACTFORCE] = True\n",
+    "model.vis.map.force = 0.005\n",
+    "\n",
+    "framerate = 30  # (Hz)\n",
+    "\n",
+    "# Simulate and display video.\n",
+    "frames = []\n",
+    "data.qpos[0] = 0.\n",
+    "data.qpos[1] = 0.\n",
+    "data.qpos[2] = 0.31\n",
+    "data.qpos[3:7] = quaternion[0]\n",
+    "data.qpos[7:] = joint_angles[0]\n",
+    "data.qpos[7:] = joint_angles[0]\n",
+    "slow_down_factor = 1\n",
+    "\n",
+    "while data.time <= 1.0*final_time:\n",
+    "    time = data.time\n",
+    "    q_joints = data.qpos[7:].copy()\n",
+    "    dq_joints = data.qvel[6:].copy()\n",
+    "    \n",
+    "    control = 120*(joint_angles_t(time) - q_joints) - 5*dq_joints\n",
+    "    data.ctrl = control #joint_angles_t(time )\n",
+    "    mujoco.mj_step(model, data)\n",
+    "\n",
+    "    if len(frames) < slow_down_factor*time * framerate:\n",
+    "        renderer.update_scene(data, scene_option = scene_option)\n",
+    "        pixels = renderer.render()\n",
+    "        frames.append(pixels)\n",
+    "media.show_video(frames, fps=framerate)\n",
+    "\n"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "control_env",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.14"
+  },
+  "orig_nbformat": 4
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}