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.
72 lines
2.6 KiB
Python
72 lines
2.6 KiB
Python
# Copyright (c) 2025 Deep Robotics
|
||
# SPDX-License-Identifier: BSD 3-Clause
|
||
|
||
"""checkpoint 解析与策略加载（IsaacLab + RSL-RL）。"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
from typing import Any
|
||
|
||
import gymnasium as gym
|
||
|
||
from isaaclab.envs import (
|
||
DirectMARLEnv,
|
||
DirectMARLEnvCfg,
|
||
DirectRLEnvCfg,
|
||
ManagerBasedRLEnvCfg,
|
||
multi_agent_to_single_agent,
|
||
)
|
||
from isaaclab_tasks.utils import get_checkpoint_path
|
||
from isaaclab_rl.rsl_rl import RslRlVecEnvWrapper
|
||
|
||
|
||
def resolve_checkpoint(
|
||
log_root_path: str,
|
||
load_run: str | None = None,
|
||
checkpoint: str | None = None,
|
||
) -> str:
|
||
"""解析出 checkpoint 的绝对路径。
|
||
|
||
允许 ``load_run`` / ``checkpoint`` 传绝对路径或带目录的相对路径，从而跨实验目录
|
||
取权重（例如从 Stage 1 的目录里取权重来评估），否则回退到 IsaacLab 的按名查找。
|
||
"""
|
||
if checkpoint is not None:
|
||
checkpoint_path = os.path.expanduser(checkpoint)
|
||
if os.path.isabs(checkpoint_path) or os.path.dirname(checkpoint_path):
|
||
checkpoint_path = os.path.abspath(checkpoint_path)
|
||
if not os.path.isfile(checkpoint_path):
|
||
raise FileNotFoundError(f"找不到 checkpoint：{checkpoint_path}")
|
||
return checkpoint_path
|
||
|
||
if load_run is not None:
|
||
run_path = os.path.expanduser(load_run)
|
||
if os.path.isabs(run_path) or os.path.dirname(run_path):
|
||
run_path = os.path.abspath(run_path)
|
||
if not os.path.isdir(run_path):
|
||
raise FileNotFoundError(f"找不到运行目录：{run_path}")
|
||
return get_checkpoint_path(os.path.dirname(run_path), os.path.basename(run_path), checkpoint)
|
||
|
||
return get_checkpoint_path(log_root_path, load_run, checkpoint)
|
||
|
||
|
||
def make_env(
|
||
task: str,
|
||
env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg,
|
||
clip_actions: float | None = None,
|
||
):
|
||
"""按已配置好的 ``env_cfg`` 建环境，并包成 RSL-RL 的向量环境。"""
|
||
env = gym.make(task, cfg=env_cfg)
|
||
if isinstance(env.unwrapped, DirectMARLEnv):
|
||
env = multi_agent_to_single_agent(env)
|
||
return RslRlVecEnvWrapper(env, clip_actions=clip_actions)
|
||
|
||
|
||
def load_policy(env, agent_cfg: Any, resume_path: str, device: str | None = None):
|
||
"""加载 checkpoint 并返回推理策略。"""
|
||
from rsl_rl.runners import OnPolicyRunner
|
||
|
||
runner = OnPolicyRunner(env, agent_cfg.to_dict(), log_dir=None, device=device or agent_cfg.device)
|
||
runner.load(resume_path)
|
||
return runner.get_inference_policy(device=device or env.unwrapped.device)
|