Go2Py_SIM/examples/06-CaT-RL-controller.ipynb

430 lines
1.4 MiB
Plaintext
Raw Normal View History

{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"RL policy based on the [SoloParkour: Constrained Reinforcement Learning for Visual Locomotion from Privileged Experience](https://arxiv.org/abs/2409.13678). "
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Flat Ground"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Test In Simulation"
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
2024-10-03 03:09:24 +08:00
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"pygame 2.6.1 (SDL 2.28.4, Python 3.10.12)\n",
"Hello from the pygame community. https://www.pygame.org/contribute.html\n"
2024-10-03 03:09:24 +08:00
]
}
],
"source": [
"from Go2Py.robot.fsm import FSM\n",
"from Go2Py.robot.remote import KeyboardRemote, XBoxRemote\n",
"from Go2Py.robot.safety import SafetyHypervisor\n",
"from Go2Py.sim.mujoco import Go2Sim\n",
"from Go2Py.control.cat import *\n",
"import torch"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"robot = Go2Sim(dt = 0.002)"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
"remote = KeyboardRemote() # XBoxRemote() # KeyboardRemote()\n",
"robot.sitDownReset()\n",
"safety_hypervisor = SafetyHypervisor(robot)"
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [],
"source": [
"class CaTController:\n",
" def __init__(self, robot, remote, checkpoint):\n",
" self.remote = remote\n",
" self.robot = robot\n",
" self.policy = Policy(checkpoint)\n",
" self.command_profile = CommandInterface()\n",
" self.agent = CaTAgent(self.command_profile, self.robot)\n",
" self.init()\n",
2024-10-03 03:09:24 +08:00
" self.hist_data = {}\n",
"\n",
" def init(self):\n",
" self.obs = self.agent.reset()\n",
" self.policy_info = {}\n",
" self.command_profile.yaw_vel_cmd = 0.0\n",
" self.command_profile.x_vel_cmd = 0.0\n",
" self.command_profile.y_vel_cmd = 0.0\n",
"\n",
" def update(self, robot, remote):\n",
" commands = remote.getCommands()\n",
" self.command_profile.yaw_vel_cmd = -commands[2]\n",
" self.command_profile.x_vel_cmd = max(commands[1] * 0.5, -0.3)\n",
" self.command_profile.y_vel_cmd = -commands[0]\n",
"\n",
" action = self.policy(self.obs, self.policy_info)\n",
2024-10-03 03:09:24 +08:00
" self.obs, self.ret, self.done, self.info = self.agent.step(action)\n",
" for key, value in self.info.items():\n",
" if key in self.hist_data:\n",
" self.hist_data[key].append(value)\n",
" else:\n",
" self.hist_data[key] = [value]"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"p_gains: [20. 20. 20. 20. 20. 20. 20. 20. 20. 20. 20. 20.]\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"/home/Go2py/Go2Py/control/cat.py:94: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
" actor_sd = torch.load(checkpoint_path, map_location=\"cpu\")\n"
]
}
],
"source": [
"from Go2Py import ASSETS_PATH \n",
"import os\n",
"checkpoint_path = os.path.join(ASSETS_PATH, 'checkpoints/SoloParkour/go2_flat.pt')\n",
"controller = CaTController(robot, remote, checkpoint_path)\n",
"decimation = 10\n",
"fsm = FSM(robot, remote, safety_hypervisor, control_dT=decimation * robot.dt, user_controller_callback=controller.update)"
]
},
{
"cell_type": "code",
"execution_count": 9,
"metadata": {},
"outputs": [],
"source": [
"remote.x_vel_cmd=0.0\n",
"remote.y_vel_cmd=0.0\n",
"remote.yaw_vel_cmd = 0.0"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Pressing `u` on the keyboard will make the robot stand up. This is equivalent to the `L2+A` combo of the Go2 builtin state machine. After the the robot is on its feet, pressing `s` will hand over the control the RL policy. This action is equivalent to the `start` key of the builtin controller. When you want to stop, pressing `u` again will act similarly to the real robot and locks it in standing mode. Finally, pressing `u` again will command the robot to sit down."
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {},
"outputs": [],
"source": [
"fsm.close()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": 11,
"metadata": {},
"outputs": [
{
"data": {
"image/png": "iVBORw0KGgoAAAANSUhEUgAABdEAAAfGCAYAAAD4GfcAAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjkuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8hTgPZAAAACXBIWXMAAA9hAAAPYQGoP6dpAAEAAElEQVR4nOy9ebxlVXnm/+wz3vnWRBVTMQSMqIgmaICYOKJoTNTEdNvddkf52UmbFGkR28TqjtEYE8zQmqQlGI2KMgQ0CTEaBREEojJZiAwikwwF1Dzc+Z5x//44Z6397rXX3nutvfe5Z91b7/fzqU9Rl1un3rvHtZ71rOf1fN/3wTAMwzAMwzAMwzAMwzAMwzBMhNKwC2AYhmEYhmEYhmEYhmEYhmEYV2ERnWEYhmEYhmEYhmEYhmEYhmFiYBGdYRiGYRiGYRiGYRiGYRiGYWJgEZ1hGIZhGIZhGIZhGIZhGIZhYmARnWEYhmEYhmEYhmEYhmEYhmFiYBGdYRiGYRiGYRiGYRiGYRiGYWJgEZ1hGIZhGIZhGIZhGIZhGIZhYmARnWEYhmEYhmEYhmEYhmEYhmFiYBGdYRiGYRiGYRiGYRiGYRiGYWJgEZ1hGIZhGIZhGIZhGIZhGIZhYmARnWFWCZ7nGf26+eabh11qodx55534nd/5HZx55pmoVqvwPG/YJTEMwzBMJo7Ed3m328Vll12GN73pTdi6dSvGx8dx+umn46Mf/SiWl5eHXR7DMAzDWHMkvs8B4DOf+Qxe8YpXYMuWLajX6zj55JNx/vnn44knnhh2aQyzIni+7/vDLoJhmHSuuOKK0J+/+MUv4oYbbsDll18e+vprX/tabNmyZSVLGygf/vCH8ad/+qc444wzMDc3h4cffhj82GIYhmFWI0fiu3x+fh6Tk5M4++yz8cu//MvYvHkzbrvtNnzhC1/Ay1/+ctx00028QM4wDMOsKo7E9zkA/M7v/A4WFxfxwhe+EOvXr8fjjz+Oz3zmM+h0OvjhD3+IY489dtglMsxAYRGdYVYpF1xwAS655JJCBOVut4tms4mRkZECKiuWPXv2YGpqCqOjo4X+zAzDMAwzbI6Ed3mz2cT3v/99/PzP/3zo6x/5yEfwoQ99CDfccAPOPffcIVXHMAzDMPk5Et7ncezYsQMveclLcPHFF+MDH/jAsMthmIHCcS4Ms4ZYWFjA+973PmzduhX1eh3Pfe5z8Zd/+ZeRl7nnebjgggtw5ZVX4gUveAHq9Tquu+46AMADDzyAV7/61RgdHcXxxx+Pj370o/jc5z4Hz/NC27Q8z8OHP/zhSA0nnXQS3vnOd4a+dvjwYVx44YWyrlNPPRV/9md/hm63m/ozbdmyBaOjo9bHgmEYhmFWI2vtXV6r1SICOgD86q/+KgDgwQcfNDgqDMMwDLO6WGvv8zhOOukk+bkMs9apDLsAhmGKwfd9vOlNb8K3v/1tvOtd78KLX/xiXH/99Xj/+9+PZ555Bp/4xCdC33/TTTfhS1/6Ei644AJs2rQJJ510Enbv3o1XvepVaLfb+MAHPoDx8XF8+tOfziViLy4u4hWveAWeeeYZ/I//8T9wwgkn4Hvf+x62b9+OXbt24a/+6q9y/uQMwzAMszY4kt7lu3fvBgBs2rQpc10MwzAM4yJr/X1+4MABdDodPPXUU/jIRz4CAHjNa16TuS6GWTX4DMOsSrZt2+bTW/hf/uVffAD+Rz/60dD3/fqv/7rveZ7/6KOPyq8B8Eulkv/AAw+EvvfCCy/0Afh33HGH/NrevXv96elpH4D/+OOPhz7jQx/6UKSuE0880X/HO94h//zHf/zH/vj4uP/www+Hvu8DH/iAXy6X/aeeeirzz8wwDMMwq5kj8V0uOPfcc/2pqSn/0KFD1n+XYRiGYVziSHuf1+t1H4APwN+4caP/N3/zN0Z/j2FWOxznwjBrhK9//esol8v4n//zf4a+/r73vQ++7+Mb3/hG6OuveMUr8PznPz/yGWeffTZ+7ud+Tn7tqKOOwtvf/vbMdX35y1/GL/7iL2L9+vXYv3+//HXuueei0+ng1ltvzfzZDMMwDLOWOFLe5X/6p3+Kb33rW/jYxz6GdevWZa6LYRiGYVxkrb/Pv/GNb+DrX/86/u///b844YQTsLCwkLkmhllNcJwLw6wRnnzySRx77LGYnJwMff15z3ue/P+Uk08+WfsZZ511VuTrz33uczPX9cgjj+Dee+/FUUcdpf3/e/fuzfzZDMMwDLOWOBLe5ddccw3+4A/+AO9617vw27/925lrYhiGYRhXWevv81e96lUAgDe84Q1485vfjNNPPx0TExO44IILMtfGMKsBFtEZ5ghlUM06O51O6M/dbhevfe1r8Xu/93va7//pn/7pgdTBMAzDMGud1fYuv+GGG/Abv/EbeOMb34hPfepTuetkGIZhmLXAanufU0455RT8zM/8DK688koW0Zk1D4voDLNGOPHEE/Gtb30Lc3NzoRXvH//4x/L/m3zGI488Evn6Qw89FPna+vXrIx24m80mdu3aFfraKaecgvn5eZx77rkmPwbDMAzDHLGs5Xf5HXfcgV/91V/FS17yEnzpS19CpcLTEIZhGGZtspbf5zqWlpbQaDQK/UyGcRHORGeYNcIv/dIvodPp4JOf/GTo65/4xCfgeR7e8IY3GH3G7bffjjvvvFN+bd++fbjyyisj33vKKadEMtM+/elPR1a7/+N//I+47bbbcP3110c+4/Dhw2i326l1MQzDMMyRwFp9lz/44IN44xvfiJNOOglf+9rXBua4YxiGYRgXWIvv83a7jUOHDkW+fuedd+K+++7DS17yktSfiWFWO2wBYZg1wq/8yq/gVa96Ff7P//k/eOKJJ/CiF70I3/zmN/GVr3wFF154IU455ZTUz/i93/s9XH755Xj961+P97znPRgfH8enP/1pnHjiibj33ntD3/vf//t/x7vf/W689a1vxWtf+1r88Ic/xPXXX49NmzaFvu/9738//vVf/xW//Mu/jHe+850488wzsbCwgPvuuw//+I//iCeeeCLydyhPPvkkLr/8cgDA97//fQDARz/6UQC91fn/9t/+m9VxYhiGYRhXWYvv8rm5OZx33nk4dOgQ3v/+9+Pf/u3fQv//lFNOwTnnnGN5pBiGYRjGXdbi+3x+fh5bt27F2972NrzgBS/A+Pg47rvvPnz+85/H9PQ0PvjBD2Y/YAyzWvAZhlmVbNu2zVdv4bm5Of+9732vf+yxx/rVatV/znOe4//FX/yF3+12Q98HwN+2bZv2c++9917/Fa94hT8yMuIfd9xx/h//8R/7n/3sZ30A/uOPPy6/r9Pp+L//+7/vb9q0yR8bG/PPO+88/9FHH/VPPPFE/x3veEekru3bt/unnnqqX6vV/E2bNvk///M/7//lX/6l32w2E3/Ob3/72z4A7a9XvOIVxseLYRiGYVzjSHiXP/7447HvcQCRf4dhGIZhVhtHwvu80Wj473nPe/wzzjjDn5qa8qvVqn/iiSf673rXu0K1MMxaxvN9319J0Z5hmNXHZZddhvPPPx+PP/44TjrppGGXwzAMwzCMJfwuZxiGYZjVD7/PGWZ4cCY6wzAMwzAMwzAMwzAMwzAMw8TAIjrDMAzDMAzDMAzDMAzDMAzDxMAiOsMwDMMwDMMwDMMwDMMwDMPEwJnoDMMwDMMwDMMwDMMwDMMwDBMDO9EZhmEYhmEYhmEYhmEYhmEYJgYW0RmGYRiGYRiGYRiGYRiGYRgmhsqwC1hJut0unn32WUxOTsLzvGGXwzAMwxzB+L6Pubk5HHvssSiVeE3bFH6XMwzDMC7B73N7+F3OMAzDuITpu/yIEtGfffZZbN26ddhlMAzDMIxk586dOP7444ddxqqB3+UMwzCMi/D73Bx+lzMMwzAukvYuP6JE9MnJSQC9gzI1NTXkahiGYZgjmdnZWWzdulW+mxgz+F3OMAzDuAS/z+3hdznDMAzjEqbv8iNKRBdbxaampvhlzTAMwzgBb2O2g9/lDMMwjIvw+9wcfpczDMMwLpL2LufQNoZhGIZhGIZhGIZhcOmll+KMM86QAvc
"text/plain": [
"<Figure size 1500x2000 with 12 Axes>"
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"import matplotlib.pyplot as plt\n",
"# Assuming 'controller.hist_data[\"torques\"]' is a dictionary with torque profiles\n",
"torques = np.array(controller.hist_data[\"torques\"])\n",
"\n",
"# Number of torque profiles\n",
"torque_nb = torques.shape[1]\n",
"\n",
"# Number of rows needed for the grid, with 3 columns per row\n",
"n_cols = 3\n",
"n_rows = int(np.ceil(torque_nb / n_cols))\n",
"\n",
"# Create the figure and axes for subplots\n",
"fig, axes = plt.subplots(n_rows, n_cols, figsize=(15, 5 * n_rows))\n",
"\n",
"# Flatten the axes array for easy indexing (in case of multiple rows)\n",
"axes = axes.flatten()\n",
"\n",
"# Plot each torque profile\n",
"for i in range(torque_nb):\n",
" axes[i].plot(torques[:, i])\n",
" axes[i].set_title(f'Torque {i+1}')\n",
" axes[i].set_xlabel('Time')\n",
" axes[i].set_ylabel('Torque Value')\n",
"\n",
"# Remove any empty subplots if torque_nb is not a multiple of 3\n",
"for j in range(torque_nb, len(axes)):\n",
" fig.delaxes(axes[j])\n",
"\n",
"# Adjust layout\n",
"plt.tight_layout()\n",
"plt.savefig(\"torque_profile.png\")\n",
"plt.show()"
]
},
{
"cell_type": "code",
"execution_count": 12,
"metadata": {},
"outputs": [
{
"data": {
"image/png": "iVBORw0KGgoAAAANSUhEUgAABdEAAAfGCAYAAAD4GfcAAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjkuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8hTgPZAAAACXBIWXMAAA9hAAAPYQGoP6dpAAEAAElEQVR4nOydeZwU1bn+n+p19mFnQBAkKoobCUaicUskosk1MYtRY2IkLok3ZMMkV3KNGJeQxWu4iUZyjWuuuZrcbP6il4goLnFBQRJRRFEEBGZgZpitZ6bX+v3RfU6d7q7q9XS9Ncz7/XzmA9PT3XOmqrrOOc95zvMapmmaYBiGYRiGYRiGYRiGYRiGYRgmDx91AxiGYRiGYRiGYRiGYRiGYRjGq7CIzjAMwzAMwzAMwzAMwzAMwzAOsIjOMAzDMAzDMAzDMAzDMAzDMA6wiM4wDMMwDMMwDMMwDMMwDMMwDrCIzjAMwzAMwzAMwzAMwzAMwzAOsIjOMAzDMAzDMAzDMAzDMAzDMA6wiM4wDMMwDMMwDMMwDMMwDMMwDrCIzjAMwzAMwzAMwzAMwzAMwzAOsIjOMAzDMAzDMAzDMAzDMAzDMA6wiM4wBwj33HMPDMPAO++8Q90U7bzzzjswDAP33HNPSc83DAPXXXddTdvEMAzDMJXA/bUF99cMwzCMV+H+2oL7a4ZJwyI6w3gM0Vm/9NJLrv3O1157Ddddd13JA4TrrrsOhmHIr4aGBsyZMwfXXHMN+vr6atvYDI888ojnOvI9e/bg6quvxoc+9CE0NzfDMAysXbuWulkMwzBMDeD+ujS82F+vWbMGX/rSl3D44YejoaEBs2bNwmWXXYY9e/ZQN41hGIbRDPfXpeHF/vqpp57Cxz/+cUyfPh11dXVoa2vDWWedhb///e/UTWNGKQHqBjAMo4cvfOELuOCCCxAOh8t+7WuvvYYf/OAHOP300zFz5sySX3f77bejqakJAwMDePTRR3HTTTfh8ccfx9///ncYhlF2O5yYMWMGhoaGEAwG5WOPPPIIbrvtNtuOfmhoCIGA+7e3LVu24Mc//jEOO+wwHHPMMXjuuedcbwPDMAzjbbi/tqDqr//t3/4N3d3dOO+883DYYYfh7bffxq233oq//vWv2LhxI9ra2lxvE8MwDOMtuL+2oOqv33jjDfh8PnzlK19BW1sb9u/fj//+7//GqaeeiocffhhnnXWW621iRjcsojPMAYLf74ff73f1d37mM5/BhAkTAABf+cpX8OlPfxp//OMf8fzzz+PEE0/U9nsMw0BdXV3Jzy/nuTqZN28eurq6MG7cOPzv//4vzjvvPJJ2MAzDMN6F+2sLqv76lltuwcknnwyfz9qUe9ZZZ+G0007DrbfeihtvvJGkXQzDMIx34P7agqq/vuyyy3DZZZdlPfav//qvmDVrFlasWMEiOuM6HOfCMCOAxx9/HKeccgoaGxsxZswYfOITn8DmzZuznmOX2TZz5kz8y7/8C5555hmccMIJqKurw6xZs3DfffdlvU6IvR/60IfkFrJKYkg+/OEPAwC2bdsGAIhEIrjqqqswffp0hMNhzJ49GzfffDNM08x63erVq3HyySdjzJgxaGpqwuzZs/G9731P/jw3s+2SSy7BbbfdBgBZ294EdpltL7/8Ms4++2y0tLSgqakJZ5xxBp5//nnbY/j3v/8dS5YswcSJE9HY2IhPfvKT2LdvX9G/v7m5GePGjSvtYDEMwzAHHNxfj4z++tRTT80S0MVj48aNyztfDMMwzIEH99cjo7+2o6GhARMnTkRPT09Fr2eYamAnOsN4nMceewxnn302Zs2aheuuuw5DQ0P4xS9+gQ9+8IPYsGFD0e1hW7duxWc+8xlceuml+OIXv4i77roLl1xyCebNm4ejjjoKp556Kr7+9a/j5z//Ob73ve/hyCOPBAD5bzm89dZbAIDx48fDNE18/OMfxxNPPIFLL70Uc+fOxd/+9jd85zvfwa5du/Czn/0MAPDqq6/iX/7lX3Dsscfi+uuvRzgcxtatWwvmnH35y1/G7t27sXr1avzmN78p2q5XX30Vp5xyClpaWvDd734XwWAQv/rVr3D66afjySefxPz587Oe/7WvfQ1jx47FsmXL8M4772DFihVYvHgxHnzwwbKPCcMwDDM64P46n5HUXw8MDGBgYEA6ABmGYZgDE+6v8/F6f93X14dYLIbOzk7cd9992LRpU9aiAMO4hskwjKe4++67TQDmiy++aJqmac6dO9ecNGmS2dXVJZ/zj3/8w/T5fObFF1+c97pt27bJx2bMmGECMJ966in52N69e81wOGxeddVV8rHf//73JgDziSeeKKmNy5YtMwGYW7ZsMfft22du27bN/NWvfmWGw2Fz8uTJZiQSMf/85z+bAMwbb7wx67Wf+cxnTMMwzK1bt5qmaZo/+9nPTADmvn37HH/ftm3bTADm3XffLR/76le/ajrdwgCYy5Ytk9+fe+65ZigUMt966y352O7du83m5mbz1FNPlY+JY7hgwQIzlUrJx7/1rW+Zfr/f7OnpKen4mGb5x5RhGIYZWXB/nc9I7K8FN9xwgwnAXLNmTdmvZRiGYbwL99f5jLT+euHChSYAE4AZCoXML3/5y+bQ0FBJr2UYnXCcC8N4mD179mDjxo245JJLsmJCjj32WHzkIx/BI488UvQ95syZg1NOOUV+P3HiRMyePRtvv/121e2bPXs2Jk6ciEMOOQRf/vKXceihh+Lhhx9GQ0MDHnnkEfj9fnz961/Pes1VV10F0zTxf//3fwCAMWPGAAD+8pe/IJVKVd2mXJLJJB599FGce+65mDVrlnx8ypQp+NznPodnnnkmr+L5FVdckbV97ZRTTkEymcT27du1t49hGIYZ+XB/XT2U/fVTTz2FH/zgB/jsZz8rt84zDMMwBx7cX1cPRX/9ox/9CI8++ijuvPNOfOADH0AsFkMikdDzBzFMGbCIzjAeRnQqs2fPzvvZkUceic7OTkQikYLvcfDBB+c9NnbsWOzfv7/q9v3hD3/A6tWrsXbtWmzduhWbNm3CvHnzZNunTp2K5ubmvHaLnwPA+eefjw9+8IO47LLLMHnyZFxwwQX43e9+p63D37dvHwYHBx2PYSqVws6dO7Mezz1mY8eOBQAtx4xhGIY58OD+unqo+uvXX38dn/zkJ3H00Ufj17/+dQUtZxiGYUYK3F9XD0V/PXfuXHzkIx/Bl770JaxevRrr1q3DJZdcUtkfwDBVwJnoDHOA41RR3MwpPlIJp556atXZofX19XjqqafwxBNP4OGHH8aqVavw4IMP4sMf/jAeffRR1yuiA7U9ZgzDMAxjB/fX5VPtMdu5cyfOPPNMtLa24pFHHskTJhiGYRgmF+6vy0fnMQuFQvj4xz+OH/3oRxgaGkJ9fX21zWOYkmEnOsN4mBkzZgAAtmzZkvez119/HRMmTEBjY2PVv0fdWqWLGTNmYPfu3ejv7896/PXXX5c/F/h8Ppxxxhm45ZZb8Nprr+Gmm27C448/jieeeKLqNk+cOBENDQ2Ox9Dn82H69OklvRfDMAzD2MH99cjrr7u6unDmmWciGo3ib3/7G6ZMmaLtvRmGYRhvwv31yOuv7RgaGoJpmnnHgmFqDYvoDONhpkyZgrlz5+Lee+9FT0+PfHzTpk149NFH8dGPflTL7xEDBfV3VMtHP/pRJJNJ3HrrrVmP/+xnP4NhGDj77LMBAN3d3XmvnTt3LgAgGo06vn+pbfb7/TjzzDPxl7/8Be+88458vKOjA7/97W9x8skno6WlpYS/iGEYhmHs4f56ZPXXkUgEH/3oR7Fr1y488sgjOOyww7S8L8MwDONtuL8eWf313r178x7r6enBH/7wB0yfPh2TJk3S8nsYplQ4zoVhPM5Pf/pTnH322TjxxBNx6aWXYmhoCL/4xS/Q2tqK6667TsvvmDt3Lvx+P3784x+jt7cX4XAYH/7wh6vqlM455xx86EMfwr//+7/jnXfewXHHHYdHH30Uf/nLX/DNb34T73nPewAA119/PZ566il87GMfw4wZM7B371788pe/xLRp03DyySc
"text/plain": [
"<Figure size 1500x2000 with 12 Axes>"
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"# Extract the joint position data for the first joint over time\n",
"joint_pos = np.array(controller.hist_data[\"joint_pos\"])[:, 0]\n",
"\n",
"# Number of data points in joint_pos\n",
"n_data_points = len(joint_pos)\n",
"\n",
"# Since you're plotting only one joint, no need for multiple subplots in this case.\n",
"# But to follow the grid requirement, we'll replicate the data across multiple subplots.\n",
"# For example, let's assume you want to visualize this data 9 times in a 3x3 grid.\n",
"\n",
"n_cols = 3\n",
"n_rows = int(np.ceil(torque_nb / n_cols))\n",
"\n",
"# Create the figure and axes for subplots\n",
"fig, axes = plt.subplots(n_rows, n_cols, figsize=(15, 5 * n_rows))\n",
"\n",
"# Flatten the axes array for easy indexing (in case of multiple rows)\n",
"axes = axes.flatten()\n",
"\n",
"# Plot the same joint position data in every subplot (as per grid requirement)\n",
"for i in range(n_rows * n_cols):\n",
" axes[i].plot(joint_pos[:, i])\n",
" axes[i].set_title(f'Joint Position {i+1}')\n",
" axes[i].set_xlabel('Time')\n",
" axes[i].set_ylabel('Position Value')\n",
"\n",
"# Adjust layout\n",
"plt.tight_layout()\n",
"plt.savefig(\"joint_position_profile.png\")\n",
"plt.show()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Test on Real Robot (ToDo)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from Go2Py.robot.fsm import FSM\n",
"from Go2Py.robot.remote import UnitreeRemote\n",
"from Go2Py.robot.safety import SafetyHypervisor\n",
"from Go2Py.control.walk_these_ways import *"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from Go2Py.robot.interface import GO2Real\n",
"import numpy as np\n",
"robot = GO2Real(mode='lowlevel')"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"remote = UnitreeRemote(robot)\n",
"safety_hypervisor = SafetyHypervisor(robot)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"class walkTheseWaysController:\n",
" def __init__(self, robot, remote, checkpoint):\n",
" self.remote = remote\n",
" self.robot = robot\n",
" self.cfg = loadParameters(checkpoint)\n",
" self.policy = Policy(checkpoint)\n",
" self.command_profile = CommandInterface()\n",
" self.agent = WalkTheseWaysAgent(self.cfg, self.command_profile, self.robot)\n",
" self.agent = HistoryWrapper(self.agent)\n",
" self.init()\n",
"\n",
" def init(self):\n",
" self.obs = self.agent.reset()\n",
" self.policy_info = {}\n",
" self.command_profile.yaw_vel_cmd = 0.0\n",
" self.command_profile.x_vel_cmd = 0.0\n",
" self.command_profile.y_vel_cmd = 0.0\n",
" self.command_profile.stance_width_cmd=0.25\n",
" self.command_profile.footswing_height_cmd=0.08\n",
" self.command_profile.step_frequency_cmd = 3.0\n",
" self.command_profile.bodyHeight = 0.00\n",
"\n",
" def update(self, robot, remote):\n",
" action = self.policy(self.obs, self.policy_info)\n",
" self.obs, self.ret, self.done, self.info = self.agent.step(action)\n",
" vy = -robot.getRemoteState().lx\n",
" vx = robot.getRemoteState().ly\n",
" omega = -robot.getRemoteState().rx*2.2\n",
" self.command_profile.x_vel_cmd = vx*1.5\n",
" self.command_profile.y_vel_cmd = vy*1.5\n",
" self.command_profile.yaw_vel_cmd = omega"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"checkpoint = \"../Go2Py/assets/checkpoints/walk_these_ways/\"\n",
"controller = walkTheseWaysController(robot, remote, checkpoint)\n",
"safety_hypervisor = SafetyHypervisor(robot)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"controller.command_profile.pitch_cmd=0.0\n",
"controller.command_profile.body_height_cmd=0.0\n",
"controller.command_profile.footswing_height_cmd=0.04\n",
"controller.command_profile.roll_cmd=0.0\n",
"controller.command_profile.stance_width_cmd=0.2\n",
"controller.command_profile.x_vel_cmd=-0.2\n",
"controller.command_profile.y_vel_cmd=0.01\n",
"controller.command_profile.setGaitType(\"trotting\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"fsm = FSM(robot, remote, safety_hypervisor, user_controller_callback=controller.update)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Pressing `L2+A` to command the robot to stand up. After the the robot is on its feet, pressing `start` will hand over the control the RL policy. When you want to stop, pressing `L2+A` again will act similarly to the factory controller and locks the robot in standing mode. Finally, pressing `L2+A` again will command the robot to sit down."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"fsm.close()"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "b1-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.12"
}
},
"nbformat": 4,
"nbformat_minor": 2
}