109 lines
3.6 KiB
Python
109 lines
3.6 KiB
Python
#!/usr/bin/env python3
|
|
"""Generate standard linear stairs for MuJoCo sim2sim testing.
|
|
|
|
Each cell: flat 2m approach -> N linear steps -> flat 2m platform -> N steps down -> flat edge.
|
|
Treads are horizontal, rises are vertical (1px = 0.1m wide, acceptable for hfield).
|
|
|
|
Usage:
|
|
uv run scripts/gen_stairs_test.py --step-height 0.07 --step-depth 0.31
|
|
"""
|
|
import numpy as np, os, argparse
|
|
from PIL import Image
|
|
|
|
HS = 0.1 # horizontal scale [m/px]
|
|
VS = 0.005 # vertical scale [m/unit]
|
|
CELL_M = 8.0
|
|
PLATFORM_M = 4.0 # bigger flat platform -> fewer steps
|
|
CELL_PX = int(CELL_M / HS) # 80
|
|
PLATFORM_PX = int(PLATFORM_M / HS) # 40
|
|
BORDER_M = 2.0
|
|
BORDER_PX = int(BORDER_M / HS) # 20
|
|
NUM_CELLS = 2
|
|
TOT_PX = NUM_CELLS * CELL_PX + 2 * BORDER_PX
|
|
TOTAL_M = TOT_PX * HS
|
|
|
|
np.random.seed(42)
|
|
|
|
|
|
def make_linear_stairs(step_height_m, step_depth_m=0.31):
|
|
"""Linear stairs: flat approach -> N steps up -> flat platform -> edge.
|
|
Each step has a flat horizontal tread and (essentially) vertical rise."""
|
|
t = np.zeros((CELL_PX, CELL_PX), dtype=np.int16)
|
|
sd = int(step_depth_m / HS) # tread depth in px
|
|
sh = int(step_height_m / VS) # rise height in pixel units
|
|
|
|
# How many steps fit on each side of the platform?
|
|
avail = (CELL_PX - PLATFORM_PX) // 2
|
|
n_steps = avail // max(sd, 1)
|
|
if n_steps < 1:
|
|
n_steps = 1
|
|
|
|
edge = (CELL_PX - PLATFORM_PX - n_steps * sd) // 2 # remaining flat on each side
|
|
|
|
# Draw steps going UP from left (in +x direction)
|
|
# Each step: flat tread at current height, then rise to next height
|
|
x = edge
|
|
h = 0
|
|
for i in range(n_steps):
|
|
x_next = x + sd
|
|
t[:, x:x_next] = h # tread at current height
|
|
x = x_next
|
|
h += sh
|
|
|
|
# Platform (flat at max height)
|
|
plat_start = x
|
|
plat_end = plat_start + PLATFORM_PX
|
|
t[:, plat_start:plat_end] = h
|
|
|
|
# Continue stairs going DOWN on the right (optional: mirror)
|
|
x = plat_end
|
|
for i in range(n_steps):
|
|
h -= sh
|
|
x_next = x + sd
|
|
t[:, x:x_next] = h
|
|
x = x_next
|
|
|
|
return t
|
|
|
|
|
|
def main():
|
|
p = argparse.ArgumentParser()
|
|
p.add_argument("--step-height", type=float, default=0.15,
|
|
help="step rise height in metres (default 0.15)")
|
|
p.add_argument("--step-depth", type=float, default=0.50,
|
|
help="step tread depth in metres (default 0.50)")
|
|
args = p.parse_args()
|
|
|
|
print(f"Linear stairs: step_h={args.step_height:.2f}m tread={args.step_depth:.2f}m platform={PLATFORM_M:.0f}m")
|
|
|
|
hf_raw = np.zeros((TOT_PX, TOT_PX), dtype=np.int16)
|
|
for i in range(NUM_CELLS):
|
|
for j in range(NUM_CELLS):
|
|
cell = make_linear_stairs(args.step_height, args.step_depth)
|
|
y0 = BORDER_PX + i * CELL_PX
|
|
x0 = BORDER_PX + j * CELL_PX
|
|
hf_raw[y0:y0 + CELL_PX, x0:x0 + CELL_PX] = cell
|
|
|
|
hf_m = hf_raw.astype(np.float32) * VS
|
|
z_min = float(hf_m.min())
|
|
z_max = float(hf_m.max())
|
|
z_range = max(z_max - z_min, 0.001)
|
|
|
|
print(f" height: [{z_min:.3f}, {z_max:.3f}]m z_scale={z_range:.3f} max={z_max*100:.0f}cm")
|
|
|
|
png = ((hf_m - z_min) / z_range * 65535.0).astype(np.uint16)
|
|
|
|
out_dir = os.path.join(os.path.dirname(__file__), "..",
|
|
"motrix_envs", "src", "motrix_envs", "locomotion",
|
|
"go1", "xmls", "assets")
|
|
os.makedirs(out_dir, exist_ok=True)
|
|
out_path = os.path.join(out_dir, "stairs_test.png")
|
|
Image.fromarray(png).save(out_path)
|
|
print(f" saved: {out_path}")
|
|
sbase = max(z_min, 0.001)
|
|
print(f" XML: size=\"{TOTAL_M/2:.1f} {TOTAL_M/2:.1f} {z_range:.3f} {sbase:.3f}\"")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|