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": [