270 lines
9.4 KiB
Python
270 lines
9.4 KiB
Python
#!/usr/bin/env python3
|
||
"""DreamWaQ 地形可视化。
|
||
|
||
键盘：
|
||
R=重置 H=高度采样点 T=遍历出生点调试 Esc=退出
|
||
|
||
用法:
|
||
uv run scripts/view_dreamwaq.py # 金字塔地形
|
||
uv run scripts/view_dreamwaq.py --flat --num-envs 1
|
||
"""
|
||
import argparse, os, sys, time
|
||
import numpy as np
|
||
|
||
os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false")
|
||
os.environ.setdefault("JAX_PLATFORMS", "cpu")
|
||
|
||
if "--flat" in sys.argv:
|
||
os.environ["DREAMWAQ_TERRAIN"] = "flat"
|
||
elif "--flat-stairs" in sys.argv:
|
||
os.environ["DREAMWAQ_TERRAIN"] = "flat_stairs"
|
||
elif "--stairs" in sys.argv:
|
||
os.environ["DREAMWAQ_TERRAIN"] = "stairs"
|
||
else:
|
||
os.environ.setdefault("DREAMWAQ_TERRAIN", "pyramid")
|
||
|
||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||
|
||
import motrix_envs.locomotion.go1.dreamwaq # noqa: F401
|
||
from motrix_envs import registry as env_registry
|
||
from motrix_envs.np.renderer import NpRenderer
|
||
from motrix_envs.math import quaternion
|
||
from motrixsim.render import RenderClosedError
|
||
|
||
# 地形类型名称（与 gen_dreamwaq_terrain.py 的 PROPORTIONS 对应）
|
||
PROPORTIONS = [0.1, 0.1, 0.35, 0.35, 0.1]
|
||
CUM = [sum(PROPORTIONS[:i + 1]) for i in range(len(PROPORTIONS))]
|
||
TYPE_NAMES = ["平滑斜坡", "粗糙斜坡", "下行楼梯", "上行楼梯", "离散障碍"]
|
||
|
||
|
||
def _cell_origin(row, col, border_m=5.0, cell_m=8.0, num_rows=10, num_cols=20):
|
||
"""计算 cell (row, col) 的中心世界坐标。"""
|
||
half_x = border_m + num_cols * cell_m / 2.0
|
||
half_y = border_m + num_rows * cell_m / 2.0
|
||
cx = -half_x + border_m + col * cell_m + cell_m / 2
|
||
cy = half_y - border_m - row * cell_m - cell_m / 2
|
||
return cx, cy
|
||
|
||
|
||
def _type_name(col):
|
||
"""根据列索引返回地形类型名称。"""
|
||
choice = col / 20 + 0.001
|
||
for i, cum in enumerate(CUM):
|
||
if choice < cum:
|
||
return TYPE_NAMES[i]
|
||
return TYPE_NAMES[-1]
|
||
|
||
|
||
def _spawn_at(env, cx, cy, spawn_z=None):
|
||
"""在指定世界坐标 spawn 单个机器人。"""
|
||
state = env._state
|
||
data = state.data
|
||
init_pos = env._init_dof_pos.copy().reshape(1, -1)
|
||
init_pos[0, 0] = cx
|
||
init_pos[0, 1] = cy
|
||
# 计算地形高度
|
||
terrain_z = float(env._sample_terrain_height(
|
||
np.array([[cx, cy]], dtype=np.float32), radius=0.35)[0])
|
||
if spawn_z is None:
|
||
spawn_z = terrain_z + 0.45 # 默认 clearance
|
||
init_pos[0, 2] = spawn_z
|
||
data.reset(env._model)
|
||
data.set_dof_pos(init_pos, env._model)
|
||
env._model.forward_kinematic(data)
|
||
# 重置 info
|
||
state.info["commands"][0] = np.array([0.0, 0.0, 0.0], dtype=np.float32)
|
||
state.info["steps"][0] = 0
|
||
state.info["obs_history"][0] = 0.0
|
||
return terrain_z, spawn_z
|
||
|
||
|
||
def main():
|
||
p = argparse.ArgumentParser(description="DreamWaQ 地形可视化")
|
||
p.add_argument("--num-envs", type=int, default=1)
|
||
p.add_argument("--flat", action="store_true")
|
||
p.add_argument("--flat-stairs", action="store_true")
|
||
p.add_argument("--stairs", action="store_true")
|
||
p.add_argument("--level", type=int, default=None)
|
||
p.add_argument("--no-stand", action="store_true")
|
||
p.add_argument("--vx", type=float, default=0.5)
|
||
args = p.parse_args()
|
||
|
||
env = env_registry.make("go1-dreamwaq-walk", num_envs=max(args.num_envs, 1))
|
||
if args.level is not None:
|
||
env._force_level = args.level
|
||
|
||
env.init_state()
|
||
n = env._num_envs
|
||
cmd = np.array([args.vx, 0.0, 0.0], dtype=np.float32)
|
||
|
||
try:
|
||
renderer = NpRenderer(env)
|
||
except Exception as e:
|
||
print(f"[ERROR] 渲染器创建失败: {e}")
|
||
renderer = None
|
||
|
||
terrain_name = os.environ.get("DREAMWAQ_TERRAIN", "pyramid")
|
||
print(f"[View] {n} 机器人 | 地形={terrain_name}")
|
||
print(f"[View] R=重置 H=高度点 T=遍历出生点 Esc=退出")
|
||
|
||
show_heights = False
|
||
traverse_mode = False
|
||
traverse_row = 0
|
||
traverse_col = 0
|
||
traverse_pending = False # 刚切换 cell, 等待稳定
|
||
traverse_settle = 0
|
||
step_count = 0
|
||
|
||
# 地形信息
|
||
num_rows = env._num_rows
|
||
num_cols = env._num_cols
|
||
cell_m = env._cell_size
|
||
|
||
def enter_traverse():
|
||
nonlocal traverse_mode, traverse_row, traverse_col, traverse_pending, traverse_settle
|
||
traverse_mode = True
|
||
traverse_row = 0
|
||
traverse_col = 0
|
||
traverse_pending = True
|
||
traverse_settle = 0
|
||
print(f"\n[T] 遍历模式: {num_rows}行 × {num_cols}列")
|
||
print(f"[T] 按 T 前进, R 退出遍历\n")
|
||
|
||
def exit_traverse():
|
||
nonlocal traverse_mode, traverse_pending
|
||
traverse_mode = False
|
||
traverse_pending = False
|
||
env.init_state()
|
||
print("[T] 退出遍历模式\n")
|
||
|
||
def advance_traverse():
|
||
nonlocal traverse_row, traverse_col, traverse_pending, traverse_settle
|
||
traverse_col += 1
|
||
if traverse_col >= num_cols:
|
||
traverse_col = 0
|
||
traverse_row += 1
|
||
if traverse_row >= num_rows:
|
||
print("[T] 遍历完成! 按 R 退出")
|
||
traverse_row = num_rows - 1
|
||
traverse_col = num_cols - 1
|
||
return False
|
||
traverse_pending = True
|
||
traverse_settle = 0
|
||
return True
|
||
|
||
def do_traverse_spawn():
|
||
"""在当前位置 spawn 并打印信息。"""
|
||
nonlocal traverse_settle
|
||
cx, cy = _cell_origin(traverse_row, traverse_col)
|
||
tname = _type_name(traverse_col)
|
||
terrain_z, spawn_z = _spawn_at(env, cx, cy)
|
||
|
||
# 标记 spawn 点（绿色球）
|
||
if renderer is not None:
|
||
g = renderer._render.gizmos
|
||
g.draw_sphere(0.15, (np.float32(cx), np.float32(cy),
|
||
np.float32(spawn_z)))
|
||
|
||
# 让机器人稳定几步
|
||
for _ in range(30):
|
||
env.step(np.zeros((1, 12), dtype=np.float32))
|
||
if renderer is not None:
|
||
renderer.render()
|
||
time.sleep(0.005)
|
||
|
||
base_z = env._body.get_pose(env._state.data)[0, 2]
|
||
contacts = env._state.info.get("contacts", np.zeros(4))
|
||
cf = env._state.info.get("privileged_obs",
|
||
np.zeros((1, 247)))[0, 45:57]
|
||
total_cf = np.sum(np.abs(cf))
|
||
|
||
print(f" r{traverse_row}c{traverse_col:02d} {tname:6s} "
|
||
f"origin=({cx:+.0f},{cy:+.0f}) "
|
||
f"terrain_z={terrain_z:.3f} spawn_z={spawn_z:.3f} "
|
||
f"base_z={base_z:.3f} cf={total_cf:.0f}N "
|
||
f"feet={contacts.astype(int).tolist()}")
|
||
|
||
traverse_settle = 30
|
||
|
||
try:
|
||
while True:
|
||
# ── 键盘 ──
|
||
if renderer is not None:
|
||
try:
|
||
inp = renderer._render.input
|
||
if inp.is_key_just_pressed("r"):
|
||
if traverse_mode:
|
||
exit_traverse()
|
||
else:
|
||
env.init_state()
|
||
step_count = 0
|
||
print("[R] 重置")
|
||
if inp.is_key_just_pressed("h"):
|
||
show_heights = not show_heights
|
||
print(f"[H] 高度点: {'开' if show_heights else '关'}")
|
||
if inp.is_key_just_pressed("t"):
|
||
if traverse_mode:
|
||
advance_traverse()
|
||
else:
|
||
enter_traverse()
|
||
except Exception:
|
||
pass
|
||
|
||
# ── 遍历模式 ──
|
||
if traverse_mode and traverse_pending:
|
||
do_traverse_spawn()
|
||
traverse_pending = False
|
||
|
||
# ── 正常模式动作 ──
|
||
if not traverse_mode or traverse_settle > 0:
|
||
if args.no_stand:
|
||
env._state.info["commands"][:] = cmd
|
||
else:
|
||
act = 0.3 * (np.random.rand(n, 12).astype(np.float32) - 0.5)
|
||
env.step(act)
|
||
if traverse_settle > 0:
|
||
traverse_settle -= 1
|
||
|
||
if step_count % 100 == 0 and step_count > 0:
|
||
bz = env._body.get_pose(env._state.data)[0, 2]
|
||
print(f"[{step_count}] base_z={bz:.3f}")
|
||
|
||
# ── 高度点 ──
|
||
if show_heights and renderer is not None:
|
||
state = env._state
|
||
pose = env._body.get_pose(state.data)
|
||
bp = pose[0, :3]
|
||
yaw = quaternion.get_yaw(pose[0:1, 3:7])[0]
|
||
cos_y, sin_y = np.cos(yaw), np.sin(yaw)
|
||
for gy in env._hy:
|
||
for gx in env._hx:
|
||
wx = bp[0] + cos_y * gx - sin_y * gy
|
||
wy = bp[1] + sin_y * gx + cos_y * gy
|
||
try:
|
||
wz = float(env._sample_terrain_height(
|
||
np.array([[wx, wy]]))[0])
|
||
g = renderer._render.gizmos
|
||
g.draw_sphere(0.02, (np.float32(wx), np.float32(wy),
|
||
np.float32(wz)))
|
||
except Exception:
|
||
pass
|
||
|
||
if renderer is not None:
|
||
renderer.render()
|
||
time.sleep(0.01)
|
||
|
||
step_count += 1
|
||
|
||
except (KeyboardInterrupt, RenderClosedError):
|
||
pass
|
||
try:
|
||
if renderer is not None:
|
||
renderer.close()
|
||
except Exception:
|
||
pass
|
||
print("[View] 结束")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|