Added height map to the Mujoco simulator.

This commit is contained in:
Rooholla-KhorramBakht 2024-10-01 21:07:58 -04:00
parent 985405740d
commit 41df8775a0
3 changed files with 33 additions and 22 deletions

View File

@ -37,6 +37,7 @@
</default>
<asset>
<hfield name="terrain" nrow="300" ncol="300" size="5 5 1. 0.1"/>
<material name="metal" rgba=".9 .95 .95 1"/>
<material name="black" rgba="0 0 0 1"/>
<material name="white" rgba="1 1 1 1"/>
@ -64,8 +65,10 @@
</asset>
<worldbody>
<geom name="heightfield" pos="0.0 0.0 -0.1" type="hfield" hfield="terrain"/>
<geom name="floor" size="0 0 0.05" type="plane" material="groundplane"/>
<body name="base" pos="0 0 0.445" childclass="go2">
<light pos="0.0 0.0 2.5" />
<inertial pos="0.021112 0 -0.005366" quat="-0.000543471 0.713435 -0.00173769 0.700719" mass="6.921"
diaginertia="0.107027 0.0980771 0.0244531"/>
<freejoint/>

View File

@ -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()]

View File

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