96 lines
3.3 KiB
Python
96 lines
3.3 KiB
Python
#!/usr/bin/env python3
|
||
"""Generate box-geom stairs XML for MuJoCo sim2sim.
|
||
|
||
Each step is a separate box with vertical rises — much steeper than hfield.
|
||
|
||
Usage:
|
||
uv run scripts/gen_stairs_box.py # default: 10 steps × 6cm = 60cm
|
||
uv run scripts/gen_stairs_box.py --step-height 0.04 --num-steps 5
|
||
uv run scripts/gen_stairs_box.py --step-height 0.10 --num-steps 8 --step-depth 0.4
|
||
"""
|
||
import argparse, os
|
||
|
||
TPL = '''<mujoco model="go1 box stairs scene">
|
||
<include file="go1_motor_actuator.xml" />
|
||
<include file="materials.xml" />
|
||
<statistic center="0 0 0.3" extent="2" meansize="0.04" />
|
||
|
||
<visual>
|
||
<headlight diffuse="0.6 0.6 0.6" ambient="0.3 0.3 0.3" specular="0 0 0" />
|
||
<rgba haze="0.15 0.25 0.35 1" />
|
||
<global azimuth="120" elevation="-20" />
|
||
<map force="0.01" />
|
||
<scale forcewidth="0.3" contactwidth="0.5" contactheight="0.2" />
|
||
<quality shadowsize="8192" />
|
||
</visual>
|
||
|
||
<worldbody>
|
||
<light pos="0 0 4" dir="0 0 -1" directional="true" />
|
||
|
||
<geom name="floor" pos="0 0 -0.001" size="0 0 0.001" type="plane"
|
||
material="motphys-ground" contype="1" conaffinity="0" priority="0" friction="0.6" />
|
||
|
||
{steps}
|
||
|
||
<!-- Fill under stairs -->
|
||
<geom name="fill" type="box" size="{fill_sx} 10 {fill_sz}" pos="{fill_x} 0 {fill_z}" rgba="0.5 0.4 0.3 1" friction="0.8 0.3 0.3"/>
|
||
</worldbody>
|
||
</mujoco>
|
||
'''
|
||
|
||
STEP_TPL = ' <geom name="step{n}" type="box" size="{sx} 10 {sz}" pos="{x} 0 {z}" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>\n'
|
||
PLAT_TPL = ' <geom name="platform" type="box" size="{sx} 10 {sz}" pos="{x} 0 {z}" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>\n'
|
||
|
||
|
||
def main():
|
||
p = argparse.ArgumentParser()
|
||
p.add_argument("--step-height", type=float, default=0.06, help="rise per step [m]")
|
||
p.add_argument("--step-depth", type=float, default=0.30, help="tread depth per step [m]")
|
||
p.add_argument("--num-steps", type=int, default=10, help="number of steps")
|
||
p.add_argument("--box-thickness", type=float, default=0.03, help="box half-height [m]")
|
||
args = p.parse_args()
|
||
|
||
h = args.step_height
|
||
d = args.step_depth
|
||
n = args.num_steps
|
||
sz = args.box_thickness # half-height of each box
|
||
|
||
steps_xml = ""
|
||
for i in range(n):
|
||
x = i * d
|
||
z = i * h + sz # center of box = step top surface - sz
|
||
steps_xml += STEP_TPL.format(n=i, sx=d/2, sz=sz, x=x, z=z)
|
||
|
||
# Platform at top
|
||
plat_x = n * d + 0.5
|
||
plat_z = n * h + sz
|
||
steps_xml += PLAT_TPL.format(sx=0.5, sz=sz, x=plat_x, z=plat_z)
|
||
|
||
# Fill box under stairs
|
||
total_depth = n * d
|
||
total_height = n * h
|
||
fill_sx = total_depth / 2
|
||
fill_sz = total_height / 2
|
||
fill_x = total_depth / 2
|
||
fill_z = -fill_sz
|
||
|
||
out = TPL.format(steps=steps_xml.rstrip(),
|
||
fill_sx=fill_sx, fill_sz=fill_sz,
|
||
fill_x=fill_x, fill_z=fill_z)
|
||
|
||
out_dir = os.path.join(os.path.dirname(__file__), "..",
|
||
"motrix_envs", "src", "motrix_envs", "locomotion",
|
||
"go1", "xmls")
|
||
out_path = os.path.join(out_dir, "scene_stairs_box.xml")
|
||
with open(out_path, "w") as f:
|
||
f.write(out)
|
||
|
||
max_h = n * h
|
||
print(f"Generated {n} steps × {h*100:.0f}cm = {max_h*100:.0f}cm total")
|
||
print(f" step depth: {d*100:.0f}cm box thickness: {sz*200:.0f}cm")
|
||
print(f" saved: {out_path}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|