394 lines
15 KiB
Python
394 lines
15 KiB
Python
"""Script to train RL agent with RSL-RL."""
|
|
|
|
import json
|
|
import logging
|
|
import os
|
|
import sys
|
|
from dataclasses import asdict, dataclass, field
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import Literal
|
|
|
|
# 训练器作为仓库内置子集直接从 scripts/ 启动，不要求额外执行 pip install -e。
|
|
TRAINER_ROOT = Path(__file__).resolve().parents[1]
|
|
SERVICE_ROOT = TRAINER_ROOT.parent
|
|
for source_root in (TRAINER_ROOT, SERVICE_ROOT):
|
|
if str(source_root) not in sys.path:
|
|
sys.path.insert(0, str(source_root))
|
|
|
|
import tyro
|
|
import warp as wp
|
|
|
|
# mjlab 1.2.0 的 GPU CUDA-graph 检查仍访问旧公开路径 wp.context；
|
|
# Warp 1.15 已把实现移到 warp._src.context，但保留了相同 runtime 契约。
|
|
if not hasattr(wp, "context"):
|
|
from warp._src import context as warp_context
|
|
|
|
wp.context = warp_context # type: ignore[attr-defined]
|
|
|
|
from mjlab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg
|
|
from mjlab.rl import MjlabOnPolicyRunner, RslRlBaseRunnerCfg, RslRlVecEnvWrapper
|
|
from mjlab.tasks.registry import list_tasks, load_env_cfg, load_rl_cfg, load_runner_cls
|
|
from mjlab.tasks.tracking.mdp import MotionCommandCfg
|
|
from mjlab.utils.gpu import select_gpus
|
|
from mjlab.utils.os import dump_yaml, get_checkpoint_path
|
|
from mjlab.utils.torch import configure_torch_backends
|
|
from mjlab.utils.wrappers import VideoRecorder
|
|
from task_config import OBSTACLE_TASK, deployment_metadata, validate_task_config
|
|
from tuning.schema import apply_reward_configuration, validate_configuration
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class TrainConfig:
|
|
env: ManagerBasedRlEnvCfg
|
|
agent: RslRlBaseRunnerCfg
|
|
motion_file: str | None = None
|
|
video: bool = False
|
|
video_length: int = 200
|
|
video_interval: int = 2000
|
|
enable_nan_guard: bool = False
|
|
torchrunx_log_dir: str | None = None
|
|
gpu_ids: list[int] | Literal["all"] | None = field(default_factory=lambda: [0])
|
|
output_dir: str | None = None
|
|
resume_checkpoint: str | None = None
|
|
pretrained_checkpoint: str | None = None
|
|
pretrained_upload_manifest: str | None = None
|
|
pretrained_onnx: str | None = None
|
|
pretrained_source_id: str | None = None
|
|
pretrained_allowed_roots: list[str] = field(default_factory=list)
|
|
reward_config: str | None = None
|
|
reward_config_json: str | None = None
|
|
task_config: str | None = None
|
|
|
|
@staticmethod
|
|
def from_task(task_id: str) -> "TrainConfig":
|
|
env_cfg = load_env_cfg(task_id)
|
|
agent_cfg = load_rl_cfg(task_id)
|
|
return TrainConfig(env=env_cfg, agent=agent_cfg)
|
|
|
|
|
|
def _load_reward_config(path: str | None, inline: str | None, task_id="Unitree-Go2-Flat") -> dict | None:
|
|
if path is not None and inline is not None:
|
|
raise ValueError("Use only one of reward_config and reward_config_json")
|
|
if inline is not None:
|
|
if len(inline.encode("utf-8")) > 64 * 1024:
|
|
raise ValueError("Reward configuration is larger than 64 KiB")
|
|
return validate_configuration(json.loads(inline), task_id)
|
|
if path is None:
|
|
return None
|
|
source = Path(path).expanduser().resolve(strict=True)
|
|
if source.stat().st_size > 64 * 1024:
|
|
raise ValueError("Reward configuration is larger than 64 KiB")
|
|
with source.open(encoding="utf-8") as stream:
|
|
return validate_configuration(json.load(stream), task_id)
|
|
|
|
|
|
def _load_task_config(task_id: str, path: str | None, seed: int) -> dict | None:
|
|
if path is None:
|
|
return validate_task_config(task_id, {}, seed)
|
|
source = Path(path).expanduser().resolve(strict=True)
|
|
if source.stat().st_size > 128 * 1024:
|
|
raise ValueError("Task configuration is larger than 128 KiB")
|
|
payload = json.loads(source.read_text(encoding="utf-8"))
|
|
allowed = {"terrainPreset", "terrainParams", "sensorCfg", "seed", "customTerrainBoxes"}
|
|
if not isinstance(payload, dict) or payload.keys() - allowed:
|
|
raise ValueError("Unknown task configuration fields")
|
|
if (
|
|
isinstance(payload.get("seed"), bool)
|
|
or not isinstance(payload.get("seed"), int)
|
|
or payload.get("seed") != seed
|
|
):
|
|
raise ValueError("Task configuration seed must equal the agent seed")
|
|
if payload.get("sensorCfg") is None:
|
|
payload.pop("sensorCfg", None)
|
|
return validate_task_config(task_id, payload, seed)
|
|
|
|
|
|
def _configure_task_and_rewards(task_id: str, cfg: TrainConfig):
|
|
custom = _load_task_config(task_id, cfg.task_config, cfg.agent.seed)
|
|
if custom is not None:
|
|
from src.tasks.obstacle_avoidance.env_cfg import apply_obstacle_configuration
|
|
from src.tasks.obstacle_avoidance.terrain import apply_terrain_configuration
|
|
if task_id == OBSTACLE_TASK:
|
|
apply_obstacle_configuration(cfg.env, custom)
|
|
else:
|
|
apply_terrain_configuration(cfg.env, custom)
|
|
deployment = deployment_metadata(task_id, custom, cfg.agent.seed)
|
|
reward_config = _load_reward_config(cfg.reward_config, cfg.reward_config_json, task_id)
|
|
if reward_config is not None:
|
|
apply_reward_configuration(cfg.env, reward_config, task_id)
|
|
if task_id == OBSTACLE_TASK:
|
|
deployment["navigation"]["speed"] = reward_config["params"]["target_velocity"]
|
|
deployment["sensorCfg"]["avoidanceWeight"] = reward_config["weights"]["avoidance_weight"]
|
|
|
|
return deployment, reward_config
|
|
|
|
|
|
def _load_pretrained(cfg: TrainConfig):
|
|
if cfg.pretrained_checkpoint is None:
|
|
if (cfg.pretrained_onnx is not None or cfg.pretrained_allowed_roots
|
|
or cfg.pretrained_source_id or cfg.pretrained_upload_manifest):
|
|
raise ValueError("Pretrained options require --pretrained-checkpoint (.pt), not ONNX alone")
|
|
return None
|
|
if cfg.resume_checkpoint is not None or cfg.agent.resume:
|
|
raise ValueError("Pretrained warm-start and resume are mutually exclusive")
|
|
from pretrained import read_pretrained_source
|
|
|
|
options = {"onnx_path": cfg.pretrained_onnx}
|
|
if cfg.pretrained_upload_manifest is not None:
|
|
if cfg.pretrained_onnx is not None:
|
|
raise ValueError("Uploaded actor cannot use adjacent ONNX sidecars")
|
|
from pretrained_upload import read_uploaded_source
|
|
read_pretrained_source = read_uploaded_source
|
|
options = {"manifest_path": cfg.pretrained_upload_manifest}
|
|
source = read_pretrained_source(
|
|
cfg.pretrained_checkpoint,
|
|
allowed_roots=cfg.pretrained_allowed_roots,
|
|
**options,
|
|
target_env=asdict(cfg.env),
|
|
target_agent=asdict(cfg.agent),
|
|
)
|
|
if cfg.pretrained_source_id is not None and source.manifest["source_id"] != cfg.pretrained_source_id:
|
|
raise ValueError("基础策略SHA身份变化，拒绝初始化")
|
|
return source
|
|
|
|
|
|
def run_train(task_id: str, cfg: TrainConfig, log_dir: Path) -> None:
|
|
deployment, reward_config = _configure_task_and_rewards(task_id, cfg)
|
|
# Validate source identity/semantics before allocating a simulation or optimizer.
|
|
pretrained = _load_pretrained(cfg)
|
|
cuda_visible = os.environ.get("CUDA_VISIBLE_DEVICES", "")
|
|
if cuda_visible == "":
|
|
device = "cpu"
|
|
seed = cfg.agent.seed
|
|
rank = 0
|
|
else:
|
|
local_rank = int(os.environ.get("LOCAL_RANK", "0"))
|
|
rank = int(os.environ.get("RANK", "0"))
|
|
# Set EGL device to match the CUDA device.
|
|
os.environ["MUJOCO_EGL_DEVICE_ID"] = str(local_rank)
|
|
device = f"cuda:{local_rank}"
|
|
# Set seed to have diversity in different processes.
|
|
seed = cfg.agent.seed + local_rank
|
|
|
|
configure_torch_backends()
|
|
|
|
cfg.agent.seed = seed
|
|
cfg.env.seed = seed
|
|
|
|
print(f"[INFO] Training with: device={device}, seed={seed}, rank={rank}")
|
|
|
|
# Check if this is a tracking task by checking for motion command.
|
|
is_tracking_task = "motion" in cfg.env.commands and isinstance(
|
|
cfg.env.commands["motion"], MotionCommandCfg
|
|
)
|
|
|
|
if is_tracking_task:
|
|
if not cfg.motion_file:
|
|
raise ValueError("For tracking tasks, --motion-file must be set ...")
|
|
motion_path = Path(cfg.motion_file).expanduser().resolve()
|
|
if not motion_path.exists():
|
|
raise FileNotFoundError(f"Motion file not found: {motion_path}")
|
|
motion_cmd = cfg.env.commands["motion"]
|
|
assert isinstance(motion_cmd, MotionCommandCfg)
|
|
motion_cmd.motion_file = str(motion_path)
|
|
print(f"[INFO] Using motion file: {motion_cmd.motion_file}")
|
|
|
|
# Check if motion_file is already set (e.g., via CLI --env.commands.motion.motion-file).
|
|
if motion_cmd.motion_file and Path(motion_cmd.motion_file).exists():
|
|
print(f"[INFO] Using local motion file: {motion_cmd.motion_file}")
|
|
|
|
# Enable NaN guard if requested.
|
|
if cfg.enable_nan_guard:
|
|
cfg.env.sim.nan_guard.enabled = True
|
|
print(f"[INFO] NaN guard enabled, output dir: {cfg.env.sim.nan_guard.output_dir}")
|
|
|
|
if rank == 0:
|
|
print(f"[INFO] Logging experiment in directory: {log_dir}")
|
|
|
|
env = ManagerBasedRlEnv(
|
|
cfg=cfg.env, device=device, render_mode="rgb_array" if cfg.video else None
|
|
)
|
|
|
|
if pretrained is not None:
|
|
from pretrained import validate_runtime_contract
|
|
|
|
try:
|
|
validate_runtime_contract(env)
|
|
except Exception:
|
|
env.close()
|
|
raise
|
|
|
|
# The ONNX runner exports this exact snapshot beside and inside policy.onnx.
|
|
env.platform_deployment = deployment
|
|
log_root_path = log_dir.parent # Go up from specific run dir to experiment dir.
|
|
|
|
resume_path: Path | None = None
|
|
explicit_resume = cfg.resume_checkpoint is not None
|
|
if explicit_resume:
|
|
resume_path = Path(cfg.resume_checkpoint).expanduser().resolve(strict=True)
|
|
elif cfg.agent.resume:
|
|
# Load checkpoint from local filesystem.
|
|
resume_path = get_checkpoint_path(
|
|
log_root_path, cfg.agent.load_run, cfg.agent.load_checkpoint
|
|
)
|
|
|
|
# Only record videos on rank 0 to avoid multiple workers writing to the same files.
|
|
if cfg.video and rank == 0:
|
|
env = VideoRecorder(
|
|
env,
|
|
video_folder=Path(log_dir) / "videos" / "train",
|
|
step_trigger=lambda step: step % cfg.video_interval == 0,
|
|
video_length=cfg.video_length,
|
|
disable_logger=True,
|
|
)
|
|
print("[INFO] Recording videos during training.")
|
|
|
|
env = RslRlVecEnvWrapper(env, clip_actions=cfg.agent.clip_actions)
|
|
|
|
agent_cfg = asdict(cfg.agent)
|
|
env_cfg = asdict(cfg.env)
|
|
|
|
runner_cls = load_runner_cls(task_id)
|
|
if runner_cls is None:
|
|
runner_cls = MjlabOnPolicyRunner
|
|
|
|
runner_kwargs = {}
|
|
runner = runner_cls(env, agent_cfg, str(log_dir), device, **runner_kwargs)
|
|
|
|
runner.add_git_repo_to_log(__file__)
|
|
if pretrained is not None:
|
|
from pretrained import initialize_runner
|
|
|
|
initialization = initialize_runner(runner, pretrained)
|
|
env.unwrapped.platform_initialization = initialization
|
|
if rank == 0:
|
|
(log_dir / "initialization.json").write_text(
|
|
json.dumps(initialization, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
|
|
)
|
|
print(f"[INFO] Warm-start actor from source {initialization['source_id']}; fresh critic/optimizer, iteration=0")
|
|
if resume_path is not None:
|
|
print(f"[INFO]: Loading model checkpoint from: {resume_path}")
|
|
runner.load(str(resume_path), map_location=device)
|
|
# Promotion copies the origin manifest, never the original actor weights.
|
|
origin_path = resume_path.parent / "initialization.json"
|
|
if origin_path.is_file():
|
|
if origin_path.stat().st_size > 64 * 1024:
|
|
raise ValueError("Initialization provenance exceeds 64 KiB")
|
|
origin = json.loads(origin_path.read_text(encoding="utf-8"))
|
|
if not isinstance(origin, dict) or origin.get("mode") != "pretrained-warm-start":
|
|
raise ValueError("Invalid initialization provenance")
|
|
env.unwrapped.platform_initialization = origin
|
|
if rank == 0:
|
|
(log_dir / "initialization.json").write_text(json.dumps(origin, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
|
|
if explicit_resume:
|
|
# RSL-RL stores the last completed zero-based iteration and otherwise
|
|
# repeats it after load. Explicit tuning promotion uses an absolute target.
|
|
runner.current_learning_iteration += 1
|
|
|
|
# Only write config files from rank 0 to avoid race conditions.
|
|
if rank == 0:
|
|
dump_yaml(log_dir / "params" / "env.yaml", env_cfg)
|
|
dump_yaml(log_dir / "params" / "agent.yaml", agent_cfg)
|
|
if reward_config is not None:
|
|
reward_snapshot = log_dir / "params" / "reward_config.json"
|
|
reward_snapshot.parent.mkdir(parents=True, exist_ok=True)
|
|
reward_snapshot.write_text(
|
|
json.dumps(reward_config, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
|
|
)
|
|
|
|
iterations = cfg.agent.max_iterations
|
|
if explicit_resume:
|
|
iterations = max(0, cfg.agent.max_iterations - runner.current_learning_iteration)
|
|
print(
|
|
f"[INFO] Learning target: current={runner.current_learning_iteration}, "
|
|
f"additional={iterations}, target={cfg.agent.max_iterations}",
|
|
flush=True,
|
|
)
|
|
runner.learn(num_learning_iterations=iterations, init_at_random_ep_len=True)
|
|
|
|
env.close()
|
|
|
|
|
|
def launch_training(task_id: str, args: TrainConfig | None = None):
|
|
args = args or TrainConfig.from_task(task_id)
|
|
|
|
# Create log directory once before launching workers.
|
|
if args.output_dir:
|
|
log_dir = Path(args.output_dir).expanduser().resolve()
|
|
else:
|
|
log_root_path = (Path("logs") / "rsl_rl" / args.agent.experiment_name).resolve()
|
|
log_dir_name = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
|
|
if args.agent.run_name:
|
|
log_dir_name += f"_{args.agent.run_name}"
|
|
log_dir = log_root_path / log_dir_name
|
|
log_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
# Select GPUs based on CUDA_VISIBLE_DEVICES and user specification.
|
|
selected_gpus, num_gpus = select_gpus(args.gpu_ids)
|
|
|
|
# Set environment variables for all modes.
|
|
if selected_gpus is None:
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = ""
|
|
else:
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(map(str, selected_gpus))
|
|
os.environ["MUJOCO_GL"] = "egl"
|
|
|
|
if num_gpus <= 1:
|
|
# CPU or single GPU: run directly without torchrunx.
|
|
run_train(task_id, args, log_dir)
|
|
else:
|
|
# Multi-GPU: use torchrunx.
|
|
import torchrunx
|
|
|
|
# torchrunx redirects stdout to logging.
|
|
logging.basicConfig(level=logging.INFO)
|
|
|
|
# Configure torchrunx logging directory.
|
|
# Priority: 1) existing env var, 2) user flag, 3) default to {log_dir}/torchrunx.
|
|
if "TORCHRUNX_LOG_DIR" not in os.environ:
|
|
if args.torchrunx_log_dir is not None:
|
|
# User specified a value via flag (could be "" to disable).
|
|
os.environ["TORCHRUNX_LOG_DIR"] = args.torchrunx_log_dir
|
|
else:
|
|
# Default: put logs in training directory.
|
|
os.environ["TORCHRUNX_LOG_DIR"] = str(log_dir / "torchrunx")
|
|
|
|
print(f"[INFO] Launching training with {num_gpus} GPUs", flush=True)
|
|
torchrunx.Launcher(
|
|
hostnames=["localhost"],
|
|
workers_per_host=num_gpus,
|
|
backend=None, # Let rsl_rl handle process group initialization.
|
|
copy_env_vars=torchrunx.DEFAULT_ENV_VARS_FOR_COPY + ("MUJOCO*",),
|
|
).run(run_train, task_id, args, log_dir)
|
|
|
|
|
|
def main():
|
|
# Parse first argument to choose the task.
|
|
# Import tasks to populate the registry.
|
|
import mjlab.tasks # noqa: F401
|
|
import src.tasks # noqa: F401
|
|
|
|
all_tasks = list_tasks()
|
|
chosen_task, remaining_args = tyro.cli(
|
|
tyro.extras.literal_type_from_choices(all_tasks),
|
|
add_help=False,
|
|
return_unknown_args=True,
|
|
config=mjlab.TYRO_FLAGS,
|
|
)
|
|
|
|
args = tyro.cli(
|
|
TrainConfig,
|
|
args=remaining_args,
|
|
default=TrainConfig.from_task(chosen_task),
|
|
prog=sys.argv[0] + f" {chosen_task}",
|
|
config=mjlab.TYRO_FLAGS,
|
|
)
|
|
del remaining_args
|
|
|
|
launch_training(task_id=chosen_task, args=args)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|