From 41df8775a0a5fd3b6dda2fe5d04ac09e91e3714e Mon Sep 17 00:00:00 2001 From: Rooholla-KhorramBakht Date: Tue, 1 Oct 2024 21:07:58 -0400 Subject: [PATCH] Added height map to the Mujoco simulator. --- Go2Py/assets/mujoco/go2.xml | 3 ++ Go2Py/sim/mujoco.py | 20 +++++++++++- .../05-walk-these-ways-RL-controller.ipynb | 32 +++++++------------ 3 files changed, 33 insertions(+), 22 deletions(-) diff --git a/Go2Py/assets/mujoco/go2.xml b/Go2Py/assets/mujoco/go2.xml index 5522c1c..c4db8fa 100644 --- a/Go2Py/assets/mujoco/go2.xml +++ b/Go2Py/assets/mujoco/go2.xml @@ -37,6 +37,7 @@ + @@ -64,8 +65,10 @@ + + diff --git a/Go2Py/sim/mujoco.py b/Go2Py/sim/mujoco.py index 0ef1829..c5b34fc 100644 --- a/Go2Py/sim/mujoco.py +++ b/Go2Py/sim/mujoco.py @@ -5,6 +5,7 @@ import numpy as np from Go2Py import ASSETS_PATH import os from scipy.spatial.transform import Rotation +import cv2 pnt = np.array([-0.2, 0, 0.05]) lidar_angles = np.linspace(0.0, 2 * np.pi, 1024).reshape(-1, 1) @@ -18,7 +19,7 @@ dist = np.zeros(nray, np.float64) class Go2Sim: - def __init__(self, mode='lowlevel', render=True, dt=0.002, xml_path=None): + def __init__(self, mode='lowlevel', render=True, dt=0.002, height_map = None, xml_path=None): if xml_path is None: self.model = mujoco.MjModel.from_xml_path( @@ -26,6 +27,13 @@ class Go2Sim: ) else: self.model = mujoco.MjModel.from_xml_path(xml_path) + + if height_map is not None: + try: + self.updateHeightMap(height_map) + except: + raise Exception('Could not set height map. Are you sure the XML contains the required asset?') + self.simulated = True self.data = mujoco.MjData(self.model) self.dt = dt @@ -102,6 +110,16 @@ class Go2Sim: else: self.step = self.stepLowlevel + def updateHeightMap(self, height_map, hfield_size = (300,300), raw_deoth_to_height_ratio = 255.): + try: + map = cv2.resize(height_map, hfield_size)/raw_deoth_to_height_ratio + self.height_map = np.flip(map, axis=0).reshape(-1) + self.model.hfield_data = self.height_map + if self.render: + self.viewer.update_hfield(0) + except: + raise Exception(f'Could not load heightmap. Make sure the heigh_map is a 2D numpy array') + def reset(self): self.q_nominal = np.hstack( [self.pos0.squeeze(), self.rot0.squeeze(), self.q0.squeeze()] diff --git a/examples/05-walk-these-ways-RL-controller.ipynb b/examples/05-walk-these-ways-RL-controller.ipynb index e5cbf0c..498d7a5 100644 --- a/examples/05-walk-these-ways-RL-controller.ipynb +++ b/examples/05-walk-these-ways-RL-controller.ipynb @@ -29,6 +29,17 @@ "robot = Go2Sim()" ] }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [], + "source": [ + "map = np.zeros((1200, 1200))\n", + "map[:200, :200] = 255\n", + "robot.updateHeightMap(map)" + ] + }, { "cell_type": "code", "execution_count": 3, @@ -90,27 +101,6 @@ "text": [ "p_gains: [20. 20. 20. 20. 20. 20. 20. 20. 20. 20. 20. 20.]\n" ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "frq: 0.06860146952866389 Hz\n", - "frq: 43.78233593252539 Hz\n", - "frq: 44.41613012538123 Hz\n", - "frq: 46.552687074074896 Hz\n", - "frq: 46.25288370349132 Hz\n", - "frq: 45.19091075603633 Hz\n", - "frq: 46.16786095609198 Hz\n", - "frq: 45.455376979181345 Hz\n", - "frq: 45.34774899450763 Hz\n", - "frq: 46.05987129647932 Hz\n", - "frq: 45.999758721663504 Hz\n", - "frq: 46.80150414532633 Hz\n", - "frq: 45.81385238828631 Hz\n", - "frq: 45.29241401652179 Hz\n", - "frq: 45.742902947880424 Hz\n" - ] } ], "source": [