224 lines
8.1 KiB
Python
224 lines
8.1 KiB
Python
#!/usr/bin/env python3
|
||
"""生成 DreamWaQ 10×20 纯 hfield 地形——OpenCV 绘制。
|
||
|
||
5 种地形类型 × 10 难度，全部在单张 PNG 高度图中。
|
||
楼梯用 1px riser 近垂直面（HS=0.05 时每像素 5cm）。
|
||
|
||
用法：
|
||
uv run python3 scripts/gen_dreamwaq_terrain.py
|
||
"""
|
||
import cv2
|
||
import numpy as np
|
||
import os
|
||
import argparse
|
||
|
||
# ═══ 参数 ═══
|
||
HS = 0.05 # 水平分辨率 [m/px]
|
||
VS = 0.005 # 垂直分辨率 [m/unit]
|
||
CELL_M = 8.0
|
||
NUM_ROWS = 10
|
||
NUM_COLS = 20
|
||
BORDER_M = 5.0
|
||
PROPORTIONS = [0.1, 0.1, 0.35, 0.35, 0.1]
|
||
CUM = [sum(PROPORTIONS[:i + 1]) for i in range(len(PROPORTIONS))]
|
||
PLATFORM_M = 3.0
|
||
_SLOPE_SCALE = 0.4 # 上游原值（已验证 z_scale 上限远超 0.54）
|
||
|
||
CELL_PX = int(CELL_M / HS) # 160
|
||
BORDER_PX = int(BORDER_M / HS) # 100
|
||
PLATFORM_PX = int(PLATFORM_M / HS) # 60
|
||
TOT_ROWS_PX = NUM_ROWS * CELL_PX + 2 * BORDER_PX # 1800
|
||
TOT_COLS_PX = NUM_COLS * CELL_PX + 2 * BORDER_PX # 3400
|
||
TOTAL_X = TOT_COLS_PX * HS
|
||
TOTAL_Y = TOT_ROWS_PX * HS
|
||
|
||
|
||
# ═══ 地形绘制 ═══
|
||
|
||
def draw_slope(canvas, x0, y0, difficulty, noise=False):
|
||
"""平滑/粗糙斜坡——与上游 pyramid_sloped_terrain 对齐。
|
||
|
||
上游逻辑：先建金字塔（中心高→边缘低），再用平台边缘高度 clip 整个 terrain，
|
||
形成与周围地形齐平的平台（而非硬清零到 0）。
|
||
"""
|
||
if difficulty <= 0:
|
||
return
|
||
slope = difficulty * _SLOPE_SCALE
|
||
max_h = int(slope * (1.0 / VS) * (CELL_M / 2.0))
|
||
if max_h <= 0:
|
||
return
|
||
cx, cy = CELL_PX // 2, CELL_PX // 2
|
||
x = np.arange(0, CELL_PX)
|
||
y = np.arange(0, CELL_PX)
|
||
xx, yy = np.meshgrid(x, y, sparse=True)
|
||
xx = (cx - np.abs(cx - xx)) / cx
|
||
yy = (cy - np.abs(cy - yy)) / cy
|
||
hf = (max_h * xx.reshape(CELL_PX, 1) * yy.reshape(1, CELL_PX)).astype(np.int32)
|
||
p2 = PLATFORM_PX // 2
|
||
# 上游 clip: 取平台边缘高度作为上下界
|
||
edge_h = int(hf[cx - p2, cy - p2])
|
||
lo = min(edge_h, 0)
|
||
hi = max(edge_h, 0)
|
||
hf = np.clip(hf, lo, hi).astype(np.uint16)
|
||
if noise:
|
||
na = int(0.05 / VS)
|
||
n = np.random.randint(-na, na + 1, (CELL_PX, CELL_PX), dtype=np.int16)
|
||
# 噪声也只在平台外
|
||
n[cx - p2:cx + p2, cy - p2:cy + p2] = 0
|
||
hf = np.clip(hf.astype(np.int32) + n, 0, 65535).astype(np.uint16)
|
||
canvas[y0:y0 + CELL_PX, x0:x0 + CELL_PX] += hf
|
||
|
||
|
||
def draw_pyramid_stairs(canvas, x0, y0, difficulty, concave=False):
|
||
"""金字塔楼梯——OpenCV 同心矩形（近垂直 riser）。
|
||
|
||
每级台阶 2px 宽（10cm tread），高度缩放保持 z_scale < 0.54。
|
||
"""
|
||
if difficulty <= 0:
|
||
return
|
||
# 上游公式：step_height = 0.05 + 0.18 * difficulty [m]
|
||
step_h_m = 0.05 + 0.18 * difficulty
|
||
step_h = max(1, int(step_h_m / VS))
|
||
cx = x0 + CELL_PX // 2
|
||
cy = y0 + CELL_PX // 2
|
||
p2 = PLATFORM_PX // 2
|
||
|
||
# 上游踏面 31cm → 6px (HS=0.05), 最多约 8 级
|
||
tread_px = max(1, int(0.31 / HS))
|
||
n_steps = min(8, (CELL_PX // 2 - p2) // tread_px)
|
||
|
||
if concave:
|
||
base_h = step_h * n_steps
|
||
cv2.rectangle(canvas, (x0, y0), (x0 + CELL_PX, y0 + CELL_PX), int(base_h), -1)
|
||
for i in range(n_steps + 1):
|
||
half = p2 + (n_steps - i) * tread_px
|
||
h = int(base_h - step_h * i)
|
||
cv2.rectangle(canvas, (cx - half, cy - half), (cx + half, cy + half), h, -1)
|
||
else:
|
||
for i in range(n_steps + 1):
|
||
half = p2 + (n_steps - i) * tread_px
|
||
h = int(step_h * i)
|
||
cv2.rectangle(canvas, (cx - half, cy - half), (cx + half, cy + half), h, -1)
|
||
|
||
|
||
def draw_obstacles(canvas, x0, y0, difficulty):
|
||
"""离散障碍物（随机矩形块）。"""
|
||
if difficulty <= 0:
|
||
return
|
||
max_h = int((0.05 + 0.2 * difficulty) / VS)
|
||
if max_h <= 0:
|
||
return
|
||
p2 = PLATFORM_PX // 2
|
||
# 上游: min_size=1.0m, max_size=2.0m, 20 个矩形
|
||
min_sz = int(1.0 / HS); max_sz = int(2.0 / HS)
|
||
for _ in range(20):
|
||
w = np.random.randint(min_sz, max_sz + 1)
|
||
ln = np.random.randint(min_sz, max_sz + 1)
|
||
si = np.random.randint(0, CELL_PX - w)
|
||
sj = np.random.randint(0, CELL_PX - ln)
|
||
cv2.rectangle(canvas, (x0 + si, y0 + sj),
|
||
(x0 + si + w, y0 + sj + ln),
|
||
int(np.random.choice([max_h // 2, max_h])), -1)
|
||
cx, cy = x0 + CELL_PX // 2, y0 + CELL_PX // 2
|
||
cv2.rectangle(canvas, (cx - p2, cy - p2), (cx + p2, cy + p2), 0, -1)
|
||
|
||
|
||
# ═══ 主流程 ═══
|
||
|
||
def main():
|
||
p = argparse.ArgumentParser()
|
||
p.add_argument("--flat-only", action="store_true")
|
||
p.add_argument("--max-level", type=int, default=None)
|
||
args = p.parse_args()
|
||
|
||
max_row = NUM_ROWS if args.max_level is None else min(args.max_level + 1, NUM_ROWS)
|
||
print(f"DreamWaQ 纯 hfield ({max_row}×{NUM_COLS}) {TOT_COLS_PX}×{TOT_ROWS_PX}px")
|
||
|
||
canvas = np.zeros((TOT_ROWS_PX, TOT_COLS_PX), dtype=np.uint16)
|
||
|
||
for row in range(max_row):
|
||
difficulty = row / NUM_ROWS
|
||
for col in range(NUM_COLS):
|
||
if args.flat_only or difficulty == 0:
|
||
continue
|
||
x0 = BORDER_PX + col * CELL_PX
|
||
y0 = BORDER_PX + row * CELL_PX
|
||
choice = col / NUM_COLS + 0.001
|
||
if choice < CUM[0]:
|
||
draw_slope(canvas, x0, y0, difficulty)
|
||
elif choice < CUM[1]:
|
||
draw_slope(canvas, x0, y0, difficulty, noise=True)
|
||
elif choice < CUM[2]:
|
||
draw_pyramid_stairs(canvas, x0, y0, difficulty, concave=True)
|
||
elif choice < CUM[3]:
|
||
draw_pyramid_stairs(canvas, x0, y0, difficulty, concave=False)
|
||
else:
|
||
draw_obstacles(canvas, x0, y0, difficulty)
|
||
|
||
hf_m = canvas.astype(np.float32) * VS
|
||
z_min, z_max = float(hf_m.min()), float(hf_m.max())
|
||
z_range = max(z_max - z_min, 0.001)
|
||
print(f" 高度范围: [{z_min:.3f}, {z_max:.3f}]m z_scale={z_range:.3f}")
|
||
|
||
if z_range > 0.54:
|
||
print(f" ⚠ z_scale={z_range:.3f} > 0.54!")
|
||
|
||
out_d = os.path.join(os.path.dirname(__file__), "..",
|
||
"motrix_envs", "src", "motrix_envs",
|
||
"locomotion", "go1", "xmls", "assets")
|
||
os.makedirs(out_d, exist_ok=True)
|
||
png = ((hf_m - z_min) / z_range * 65535.0).astype(np.uint16)
|
||
cv2.imwrite(os.path.join(out_d, "dreamwaq_terrain.png"), png)
|
||
|
||
# XML
|
||
xml = f"""<mujoco model="go1 dreamwaq terrain scene">
|
||
<include file="go1_motor_actuator.xml" />
|
||
<include file="materials.xml" />
|
||
<statistic center="0 0 0.2" extent="5" 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>
|
||
|
||
<asset>
|
||
<hfield name="dreamwaq_terrain"
|
||
file="assets/dreamwaq_terrain.png"
|
||
size="{TOTAL_X / 2:.1f} {TOTAL_Y / 2:.1f} {z_range:.3f} {max(z_min, 0.001):.3f}" />
|
||
</asset>
|
||
|
||
<worldbody>
|
||
<light pos="0 0 4" dir="0 0 -1" directional="true" />
|
||
<geom name="floor" pos="0 0 0" type="hfield" hfield="dreamwaq_terrain"
|
||
material="motphys-ground" contype="1" conaffinity="0"
|
||
priority="1" friction="0.6" />
|
||
</worldbody>
|
||
|
||
<sensor>
|
||
<contact name="FR_foot_contact" geom2="FR_foot" geom1="floor" data="force" num="1" />
|
||
<contact name="FL_foot_contact" geom2="FL_foot" geom1="floor" data="force" num="1" />
|
||
<contact name="RR_foot_contact" geom2="RR_foot" geom1="floor" data="force" num="1" />
|
||
<contact name="RL_foot_contact" geom2="RL_foot" geom1="floor" data="force" num="1" />
|
||
</sensor>
|
||
</mujoco>
|
||
"""
|
||
xml_dir = os.path.join(os.path.dirname(__file__), "..",
|
||
"motrix_envs", "src", "motrix_envs",
|
||
"locomotion", "go1", "xmls")
|
||
with open(os.path.join(xml_dir, "scene_dreamwaq_terrain.xml"), "w") as f:
|
||
f.write(xml)
|
||
|
||
half_x = TOTAL_X / 2
|
||
half_y = TOTAL_Y / 2
|
||
print(f" XML: size=\"{half_x:.1f} {half_y:.1f} {z_range:.3f} {max(z_min, 0.001):.3f}\"")
|
||
print(f" 楼梯: 1px tread (5cm), 1px riser → 近垂直面")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
np.random.seed(42)
|
||
main()
|