syfffffjj
07027e8fde
Current code snapshot: smooth/bridge task configs, smooth random terrain, roll/pitch stability rewards, M20 URDF updates, eval profile and stability training documentation.
308 lines
13 KiB
Python
308 lines
13 KiB
Python
# Copyright (c) 2025 Deep Robotics
|
||
# SPDX-License-Identifier: BSD 3-Clause
|
||
|
||
"""把评估协议与地形规格翻译成 IsaacLab 的 ``env_cfg`` 改动。
|
||
|
||
这些函数对应旧 ``eval.py`` 里的 ``apply_eval_terrain_baseline`` /
|
||
``configure_eval_env`` / ``configure_upper_limit_env`` / ``assign_upper_limit_origins``，
|
||
区别是：
|
||
|
||
- 数值全部来自 :class:`~eval_kit.core.terrain_spec.TerrainSpec` 与
|
||
:class:`~eval_kit.core.schema.EvalProtocol`，没有散落的常数
|
||
- 需要关掉的随机化事件分「通用」与「机器人相关」两组，机器人相关的那组用
|
||
``hasattr`` 守卫，因此换机器人不会因为字段不存在而报错
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import copy
|
||
from typing import Any, Sequence
|
||
|
||
import torch
|
||
|
||
from isaaclab.terrains.config.rough import ROUGH_TERRAINS_CFG
|
||
|
||
from eval_kit.core.schema import EvalProtocol
|
||
from eval_kit.core.terrain_spec import TerrainSpec
|
||
|
||
#: 所有任务都需要的通用随机化项。
|
||
COMMON_RANDOMIZATION_TERMS: tuple[str, ...] = (
|
||
"randomize_rigid_body_material",
|
||
"randomize_rigid_body_mass",
|
||
"randomize_rigid_body_mass_base",
|
||
"randomize_rigid_body_inertia",
|
||
"randomize_com_positions",
|
||
"randomize_apply_external_force_torque",
|
||
"randomize_actuator_gains",
|
||
)
|
||
|
||
#: 火蚁专属的随机化项（机身/膝部与轮胎分开做材质随机化）。
|
||
ROBOT_SPECIFIC_RANDOMIZATION_TERMS: tuple[str, ...] = (
|
||
"randomize_wheel_material",
|
||
"randomize_support_material",
|
||
)
|
||
|
||
|
||
def set_fixed_command(env_cfg: Any, protocol: EvalProtocol) -> None:
|
||
"""把速度指令钉成一个恒定值，并在整个 episode 内不重采样。"""
|
||
command_cfg = env_cfg.commands.base_velocity
|
||
command_x, command_y, command_yaw = protocol.command
|
||
command_cfg.heading_command = False
|
||
command_cfg.resampling_time_range = (
|
||
protocol.steps * protocol.dt + 0.01,
|
||
protocol.steps * protocol.dt + 0.01,
|
||
)
|
||
command_cfg.ranges.lin_vel_x = (command_x, command_x)
|
||
command_cfg.ranges.lin_vel_y = (command_y, command_y)
|
||
command_cfg.ranges.ang_vel_z = (command_yaw, command_yaw)
|
||
command_cfg.debug_vis = False
|
||
|
||
|
||
def apply_eval_terrain_baseline(env_cfg: Any, spec: TerrainSpec, protocol: EvalProtocol) -> None:
|
||
"""所有任务统一用同一套评估地形，保证不同机器人/不同实验面对的是同一块地。"""
|
||
if env_cfg.scene.terrain.terrain_type == "plane":
|
||
env_cfg.scene.terrain.terrain_type = "generator"
|
||
env_cfg.scene.terrain.terrain_generator = copy.deepcopy(ROUGH_TERRAINS_CFG)
|
||
|
||
generator = env_cfg.scene.terrain.terrain_generator
|
||
generator.size = (protocol.terrain_size_m, protocol.terrain_size_m)
|
||
generator.border_width = 2.0
|
||
generator.curriculum = False
|
||
generator.num_rows = spec.grid_rows
|
||
generator.num_cols = spec.grid_cols
|
||
generator.use_cache = False
|
||
env_cfg.scene.terrain.max_init_terrain_level = None
|
||
|
||
# 瓦片变小的时候，台阶宽度/平台宽度等特征也要跟着收，否则一块地里塞不下
|
||
for sub_cfg in generator.sub_terrains.values():
|
||
for field_name in spec.feature_scaling:
|
||
if hasattr(sub_cfg, field_name):
|
||
setattr(sub_cfg, field_name, spec.scaling_for(field_name, protocol.terrain_size_m))
|
||
|
||
|
||
def _build_push_event(protocol: EvalProtocol):
|
||
"""构造周期性扰动事件。"""
|
||
from isaaclab.managers import EventTermCfg as EventTerm
|
||
|
||
try:
|
||
# 与旧 eval.py 同源，保证扰动行为完全一致
|
||
import rl_training.tasks.manager_based.locomotion.velocity.mdp as mdp
|
||
except ImportError:
|
||
from isaaclab_tasks.manager_based.locomotion.velocity import mdp
|
||
|
||
return EventTerm(
|
||
func=mdp.push_by_setting_velocity,
|
||
mode="interval",
|
||
interval_range_s=(protocol.push_interval_s[0], protocol.push_interval_s[1]),
|
||
params={
|
||
"velocity_range": {
|
||
"x": protocol.push_velocity_range,
|
||
"y": protocol.push_velocity_range,
|
||
}
|
||
},
|
||
)
|
||
|
||
|
||
def _configure_viewer(env_cfg: Any, viewer_env_index: int | None, free_camera: bool) -> None:
|
||
"""设置视口相机：跟随某个机器人，或交给客户端自己控制。
|
||
|
||
``ViewportCameraController`` 只在 ``origin_type`` 是 ``asset_root`` /
|
||
``asset_body`` 时每帧改写相机位姿；设成 ``world`` 只在初始化时写一次，
|
||
之后客户端的操作不会被抢回去 —— 这就是「自由视角」。起始视角由子进程按
|
||
实际地形范围另设（见 ``eval_suite._set_free_camera_overview``）。
|
||
"""
|
||
viewer = getattr(env_cfg, "viewer", None)
|
||
if viewer is None:
|
||
return
|
||
if free_camera:
|
||
viewer.origin_type = "world"
|
||
# 仅作初始化时的临时值，正式视角在环境建好后按地形范围重设
|
||
viewer.eye = (0.0, -20.0, 15.0)
|
||
viewer.lookat = (0.0, 0.0, 0.0)
|
||
return
|
||
if viewer_env_index is None:
|
||
return
|
||
viewer.origin_type = "asset_root"
|
||
viewer.asset_name = "robot"
|
||
viewer.env_index = int(viewer_env_index)
|
||
# 相对机器人的偏移：斜前方 2.5 m、抬高 1.6 m，看向机身略前方
|
||
viewer.eye = (2.5, -2.5, 1.6)
|
||
viewer.lookat = (0.4, 0.0, 0.1)
|
||
|
||
|
||
def configure_eval_env(
|
||
env_cfg: Any,
|
||
spec: TerrainSpec,
|
||
protocol: EvalProtocol,
|
||
*,
|
||
num_envs: int | None = None,
|
||
device: str | None = None,
|
||
seed: int | None = None,
|
||
viewer_env_index: int | None = None,
|
||
free_camera: bool = False,
|
||
) -> None:
|
||
"""把训练用配置改造成单次评估用配置。
|
||
|
||
关掉观测噪声、课程、随机化，把重置位姿钉死，指令固定，地形统一。
|
||
|
||
``viewer_env_index`` 不为 None 时，把视口相机改成跟随该环境的机器人；
|
||
``free_camera`` 为 True 时改成不跟随任何物体、由客户端自己移动（两者互斥，
|
||
前者优先）。这两种都只在有画面时（本地 GUI 或推流）才有意义 —— headless
|
||
且不推流时 IsaacLab 根本不创建相机控制器，设了也无害。
|
||
"""
|
||
env_cfg.scene.num_envs = num_envs if num_envs is not None else env_cfg.scene.num_envs
|
||
env_cfg.episode_length_s = protocol.steps * protocol.dt
|
||
if seed is not None:
|
||
env_cfg.seed = seed
|
||
if device is not None:
|
||
env_cfg.sim.device = device
|
||
|
||
_configure_viewer(env_cfg, viewer_env_index, free_camera)
|
||
|
||
env_cfg.observations.policy.enable_corruption = False
|
||
set_fixed_command(env_cfg, protocol)
|
||
apply_eval_terrain_baseline(env_cfg, spec, protocol)
|
||
|
||
# 课程必须关掉，否则难度/指令会在评估过程中自己变
|
||
for term in ("terrain_levels", "command_levels"):
|
||
if getattr(env_cfg.curriculum, term, None) is not None:
|
||
setattr(env_cfg.curriculum, term, None)
|
||
|
||
for term in (*COMMON_RANDOMIZATION_TERMS, *ROBOT_SPECIFIC_RANDOMIZATION_TERMS):
|
||
if hasattr(env_cfg.events, term):
|
||
setattr(env_cfg.events, term, None)
|
||
|
||
# 重置位姿钉死：同一批机器人从同一个位置、同一姿态、零速度出发
|
||
reset_base = getattr(env_cfg.events, "randomize_reset_base", None)
|
||
if reset_base is not None:
|
||
start_x, start_y, start_yaw = protocol.start_pose
|
||
# z 是相对地形瓦片中心高度的抬升。默认 0.0 与旧 eval.py 一致，但只对
|
||
# 「瓦片中心最高」的地形成立；倒金字塔台阶这类中心最低的地形需要正值，
|
||
# 否则机器人在 |x|>platform 处会被生成到地形几何体内部而卡住。
|
||
reset_base.params = {
|
||
"pose_range": {
|
||
"x": (start_x, start_x),
|
||
"y": (start_y, start_y),
|
||
"z": (protocol.start_z_offset_m, protocol.start_z_offset_m),
|
||
"roll": (0.0, 0.0),
|
||
"pitch": (0.0, 0.0),
|
||
"yaw": (start_yaw, start_yaw),
|
||
},
|
||
"velocity_range": {axis: (0.0, 0.0) for axis in ("x", "y", "z", "roll", "pitch", "yaw")},
|
||
}
|
||
|
||
if hasattr(env_cfg.events, "randomize_push_robot"):
|
||
env_cfg.events.randomize_push_robot = _build_push_event(protocol) if protocol.enable_interval_pushes else None
|
||
|
||
# 越界终止项：它比较的是世界坐标与整张地图的范围，不是机器人自己那块地砖，
|
||
# 小地砖 + 大网格时会把最外侧格子判成越界（详见 EvalProtocol 的说明）。
|
||
if not protocol.enforce_terrain_bounds and getattr(env_cfg.terminations, "terrain_out_of_bounds", None) is not None:
|
||
env_cfg.terminations.terrain_out_of_bounds = None
|
||
|
||
|
||
# 说明：判定地砖是否够大的逻辑已移到 eval_kit.core.schema.EvalProtocol.
|
||
# minimum_terrain_size_for_bounds()。它只用到协议自己的 start_pose，放在适配层会
|
||
# 让不启动仿真的编排进程被迫 import isaaclab（进而因缺 pxr 而失败）。
|
||
|
||
|
||
def configure_upper_limit_env(
|
||
env_cfg: Any,
|
||
spec: TerrainSpec,
|
||
protocol: EvalProtocol,
|
||
terrain_types: Sequence[str],
|
||
*,
|
||
seed: int | None = None,
|
||
device: str | None = None,
|
||
viewer_env_index: int | None = None,
|
||
free_camera: bool = False,
|
||
) -> None:
|
||
"""配置上限评估的固定网格：行 = 难度等级，列 = 地形类型。
|
||
|
||
``seed`` 同时决定地形生成与重置采样。显式传入是为了让地形本身在不同种子间
|
||
保持一致（同一块地），而机器人初始状态可以随种子变化。
|
||
|
||
``viewer_env_index`` / ``free_camera`` 透传给 :func:`configure_eval_env`，
|
||
用于让视口相机跟随某个环境的机器人，或者干脆交给客户端自由控制。
|
||
"""
|
||
configure_eval_env(
|
||
env_cfg, spec, protocol, seed=seed, device=device,
|
||
viewer_env_index=viewer_env_index, free_camera=free_camera,
|
||
)
|
||
|
||
generator = env_cfg.scene.terrain.terrain_generator
|
||
generator.curriculum = True
|
||
generator.num_rows = protocol.num_levels
|
||
generator.num_cols = len(terrain_types)
|
||
generator.difficulty_range = (0.0, 1.0)
|
||
if hasattr(generator, "seed"):
|
||
generator.seed = int(seed) if seed is not None else int(env_cfg.seed or 0)
|
||
|
||
# 每个格子里放 envs_per_cell 个机器人
|
||
env_cfg.scene.num_envs = protocol.num_levels * len(terrain_types) * protocol.envs_per_cell
|
||
env_cfg.scene.terrain.max_init_terrain_level = protocol.num_levels - 1
|
||
|
||
# 只保留这次评估要用到的子地形，且按被选中的次数分配占比
|
||
isaaclab_names = [spec.isaaclab_subterrain[terrain_type] for terrain_type in terrain_types]
|
||
original = generator.sub_terrains
|
||
generator.sub_terrains = {name: copy.deepcopy(original[name]) for name in isaaclab_names}
|
||
for name, sub_cfg in generator.sub_terrains.items():
|
||
sub_cfg.proportion = isaaclab_names.count(name) / len(isaaclab_names)
|
||
|
||
# 把等级表里的首尾值写进子地形的难度区间。
|
||
# 用 _set_sub_terrain_field 而不是裸 setattr：字段名写错要立刻报错，
|
||
# 静默造一个假属性会让地形保持默认难度，而结果看起来「只是很难」。
|
||
for terrain_type in terrain_types:
|
||
sub_cfg = generator.sub_terrains[spec.isaaclab_subterrain[terrain_type]]
|
||
levels = spec.levels_for(terrain_type, protocol.num_levels)
|
||
_set_sub_terrain_field(sub_cfg, spec.isaaclab_param[terrain_type], (levels[0], levels[-1]), terrain_type)
|
||
|
||
# 与等级无关的固定字段（例如粗糙地形的噪声步长）
|
||
for terrain_type, extra in spec.isaaclab_extra_params.items():
|
||
if terrain_type not in terrain_types:
|
||
continue
|
||
sub_cfg = generator.sub_terrains[spec.isaaclab_subterrain[terrain_type]]
|
||
for field_name, value in extra.items():
|
||
_set_sub_terrain_field(sub_cfg, field_name, value, terrain_type)
|
||
|
||
|
||
def _set_sub_terrain_field(sub_cfg: Any, field_name: str, value: Any, terrain_type: str) -> None:
|
||
"""写入子地形字段，字段不存在时报错而不是静默创建一个假属性。"""
|
||
if not hasattr(sub_cfg, field_name):
|
||
raise AttributeError(
|
||
f"地形规格为 {terrain_type!r} 指定的字段 {field_name!r} 在 "
|
||
f"{type(sub_cfg).__name__} 上不存在。裸 setattr 会静默造一个假属性、"
|
||
f"使地形保持默认难度，而那看起来只是「很难」。可用字段："
|
||
f"{sorted(k for k in vars(sub_cfg) if not k.startswith('_'))}"
|
||
)
|
||
setattr(sub_cfg, field_name, value)
|
||
|
||
|
||
def assign_upper_limit_origins(
|
||
env: Any,
|
||
spec: TerrainSpec,
|
||
protocol: EvalProtocol,
|
||
terrain_types: Sequence[str],
|
||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||
"""给每个格子平均分配机器人，并直接把地形等级/类型/起点写进地形对象。
|
||
|
||
返回 ``(level_ids, type_ids)``，用于事后把逐环境结果折叠成矩阵。
|
||
"""
|
||
base_env = env.unwrapped
|
||
terrain = base_env.scene.terrain
|
||
device = base_env.device
|
||
|
||
levels: list[int] = []
|
||
types: list[int] = []
|
||
for type_index in range(len(terrain_types)):
|
||
for level_index in range(protocol.num_levels):
|
||
levels.extend([level_index] * protocol.envs_per_cell)
|
||
types.extend([type_index] * protocol.envs_per_cell)
|
||
|
||
level_ids = torch.tensor(levels, dtype=torch.long, device=device)
|
||
type_ids = torch.tensor(types, dtype=torch.long, device=device)
|
||
terrain.terrain_levels = level_ids
|
||
terrain.terrain_types = type_ids
|
||
terrain.env_origins = terrain.terrain_origins[level_ids, type_ids]
|
||
env.reset()
|
||
return level_ids, type_ids
|