101 lines
3.5 KiB
Python
101 lines
3.5 KiB
Python
#!/usr/bin/env python3
|
||
"""Generate symmetric box-geom stairs XML for MuJoCo sim2sim.
|
||
|
||
The course is a flat approach, ascending stairs, a top platform, and matching
|
||
descending stairs. Each tread is a solid box rising from the ground.
|
||
|
||
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}
|
||
</worldbody>
|
||
</mujoco>
|
||
'''
|
||
|
||
STEP_TPL = ' <geom name="{name}" type="box" size="{sx:.6g} 10 {sz:.6g}" pos="{x:.6g} 0 {z:.6g}" 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:.6g} 10 {sz:.6g}" pos="{x:.6g} 0 {z:.6g}" 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("--platform-depth", type=float, default=1.0, help="top platform depth [m]")
|
||
args = p.parse_args()
|
||
|
||
h = args.step_height
|
||
d = args.step_depth
|
||
n = args.num_steps
|
||
platform_depth = args.platform_depth
|
||
|
||
steps_xml = ""
|
||
for i in range(n):
|
||
top = (i + 1) * h
|
||
steps_xml += STEP_TPL.format(
|
||
name=f"step_up_{i}", sx=d / 2, sz=top / 2,
|
||
x=(i + 0.5) * d, z=top / 2,
|
||
)
|
||
|
||
# Platform at top
|
||
total_height = n * h
|
||
platform_start = n * d
|
||
platform_end = platform_start + platform_depth
|
||
steps_xml += PLAT_TPL.format(
|
||
sx=platform_depth / 2, sz=total_height / 2,
|
||
x=platform_start + platform_depth / 2, z=total_height / 2,
|
||
)
|
||
|
||
for i in range(n):
|
||
top = (n - i - 1) * h
|
||
if top <= 0:
|
||
continue
|
||
steps_xml += STEP_TPL.format(
|
||
name=f"step_down_{i}", sx=d / 2, sz=top / 2,
|
||
x=platform_end + (i + 0.5) * d, z=top / 2,
|
||
)
|
||
|
||
out = TPL.format(steps=steps_xml.rstrip())
|
||
|
||
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
|
||
finish_x = platform_end + n * d
|
||
print(f"Generated {n} steps up/down x {h*100:.0f}cm = {max_h*100:.0f}cm total")
|
||
print(f" step depth: {d*100:.0f}cm platform: {platform_depth:.2f}m")
|
||
print(f" course: x=0.00m to x={finish_x:.2f}m")
|
||
print(f" saved: {out_path}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|