Add standalone Huoyi training project
This commit is contained in:
huangweifeng
@@ -0,0 +1,44 @@
|
||||
# Copyright (c) 2025 Deep Robotics
|
||||
# SPDX-License-Identifier: BSD 3-Clause
|
||||
|
||||
# Copyright (c) 2024-2025 Ziqi Fan
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import torch
|
||||
|
||||
import isaaclab.utils.math as math_utils
|
||||
|
||||
|
||||
def camera_follow(env):
|
||||
if not hasattr(camera_follow, "smooth_camera_positions"):
|
||||
camera_follow.smooth_camera_positions = []
|
||||
robot_pos = env.unwrapped.scene["robot"].data.root_pos_w[0]
|
||||
robot_quat = env.unwrapped.scene["robot"].data.root_quat_w[0]
|
||||
camera_offset = torch.tensor([-3.0, 0.0, 0.5], dtype=torch.float32, device=env.device)
|
||||
camera_pos = math_utils.transform_points(
|
||||
camera_offset.unsqueeze(0), pos=robot_pos.unsqueeze(0), quat=robot_quat.unsqueeze(0)
|
||||
).squeeze(0)
|
||||
# camera_pos[2] = torch.clamp(camera_pos[2], min=0.1)
|
||||
window_size = 50
|
||||
camera_follow.smooth_camera_positions.append(camera_pos)
|
||||
if len(camera_follow.smooth_camera_positions) > window_size:
|
||||
camera_follow.smooth_camera_positions.pop(0)
|
||||
smooth_camera_pos = torch.mean(torch.stack(camera_follow.smooth_camera_positions), dim=0)
|
||||
env.unwrapped.viewport_camera_controller.set_view_env_index(env_index=0)
|
||||
env.unwrapped.viewport_camera_controller.update_view_location(
|
||||
eye=smooth_camera_pos.cpu().numpy(), lookat=robot_pos.cpu().numpy()
|
||||
)
|
||||
|
||||
|
||||
|
||||
|
||||
import os
|
||||
import pickle
|
||||
from typing import Any
|
||||
|
||||
def reindex_feet(vec):
|
||||
return vec[:, [1, 0, 3, 2]]
|
||||
|
||||
def reindex(vec):
|
||||
return vec[:, [3, 4, 5, 0, 1, 2, 9, 10, 11, 6, 7, 8]]
|
||||
|
||||
@@ -0,0 +1,97 @@
|
||||
# Copyright (c) 2025 Deep Robotics
|
||||
# SPDX-License-Identifier: BSD 3-Clause
|
||||
|
||||
# Copyright (c) 2024-2025 Ziqi Fan
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# Copyright (c) 2024-2025, The Isaac Lab Project Developers.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import random
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from isaaclab_rl.rsl_rl import RslRlOnPolicyRunnerCfg
|
||||
|
||||
|
||||
def add_rsl_rl_args(parser: argparse.ArgumentParser):
|
||||
"""Add RSL-RL arguments to the parser.
|
||||
|
||||
Args:
|
||||
parser: The parser to add the arguments to.
|
||||
"""
|
||||
# create a new argument group
|
||||
arg_group = parser.add_argument_group("rsl_rl", description="Arguments for RSL-RL agent.")
|
||||
# -- experiment arguments
|
||||
arg_group.add_argument(
|
||||
"--experiment_name", type=str, default=None, help="Name of the experiment folder where logs will be stored."
|
||||
)
|
||||
arg_group.add_argument("--run_name", type=str, default=None, help="Run name suffix to the log directory.")
|
||||
# -- load arguments
|
||||
arg_group.add_argument("--resume", action="store_true", default=False, help="Whether to resume from a checkpoint.")
|
||||
arg_group.add_argument("--load_run", type=str, default=None, help="Name of the run folder to resume from.")
|
||||
arg_group.add_argument("--checkpoint", type=str, default=None, help="Checkpoint file to resume from.")
|
||||
# -- logger arguments
|
||||
arg_group.add_argument(
|
||||
"--logger", type=str, default=None, choices={"wandb", "tensorboard", "neptune"}, help="Logger module to use."
|
||||
)
|
||||
arg_group.add_argument(
|
||||
"--log_project_name", type=str, default=None, help="Name of the logging project when using wandb or neptune."
|
||||
)
|
||||
|
||||
|
||||
def parse_rsl_rl_cfg(task_name: str, args_cli: argparse.Namespace) -> RslRlOnPolicyRunnerCfg:
|
||||
"""Parse configuration for RSL-RL agent based on inputs.
|
||||
|
||||
Args:
|
||||
task_name: The name of the environment.
|
||||
args_cli: The command line arguments.
|
||||
|
||||
Returns:
|
||||
The parsed configuration for RSL-RL agent based on inputs.
|
||||
"""
|
||||
from isaaclab_tasks.utils.parse_cfg import load_cfg_from_registry
|
||||
|
||||
# load the default configuration
|
||||
rslrl_cfg: RslRlOnPolicyRunnerCfg = load_cfg_from_registry(task_name, "rsl_rl_cfg_entry_point")
|
||||
rslrl_cfg = update_rsl_rl_cfg(rslrl_cfg, args_cli)
|
||||
return rslrl_cfg
|
||||
|
||||
|
||||
def update_rsl_rl_cfg(agent_cfg: RslRlOnPolicyRunnerCfg, args_cli: argparse.Namespace):
|
||||
"""Update configuration for RSL-RL agent based on inputs.
|
||||
|
||||
Args:
|
||||
agent_cfg: The configuration for RSL-RL agent.
|
||||
args_cli: The command line arguments.
|
||||
|
||||
Returns:
|
||||
The updated configuration for RSL-RL agent based on inputs.
|
||||
"""
|
||||
# override the default configuration with CLI arguments
|
||||
if hasattr(args_cli, "seed") and args_cli.seed is not None:
|
||||
# randomly sample a seed if seed = -1
|
||||
if args_cli.seed == -1:
|
||||
args_cli.seed = random.randint(0, 10000)
|
||||
agent_cfg.seed = args_cli.seed
|
||||
if args_cli.resume is not None:
|
||||
agent_cfg.resume = args_cli.resume
|
||||
if args_cli.load_run is not None:
|
||||
agent_cfg.load_run = args_cli.load_run
|
||||
if args_cli.checkpoint is not None:
|
||||
agent_cfg.load_checkpoint = args_cli.checkpoint
|
||||
if args_cli.run_name is not None:
|
||||
agent_cfg.run_name = args_cli.run_name
|
||||
if args_cli.logger is not None:
|
||||
agent_cfg.logger = args_cli.logger
|
||||
# set the project name for wandb and neptune
|
||||
if agent_cfg.logger in {"wandb", "neptune"} and args_cli.log_project_name:
|
||||
agent_cfg.wandb_project = args_cli.log_project_name
|
||||
agent_cfg.neptune_project = args_cli.log_project_name
|
||||
|
||||
return agent_cfg
|
||||
@@ -0,0 +1,716 @@
|
||||
# Copyright (c) 2025 Deep Robotics
|
||||
# SPDX-License-Identifier: BSD 3-Clause
|
||||
|
||||
"""Evaluate an RSL-RL checkpoint with criterion metrics migrated from legged_gym."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import copy
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from datetime import datetime
|
||||
|
||||
import numpy as np
|
||||
|
||||
PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", ".."))
|
||||
PROJECT_SOURCE_PATH = os.path.join(PROJECT_ROOT, "source", "rl_training")
|
||||
if PROJECT_SOURCE_PATH not in sys.path:
|
||||
sys.path.insert(0, PROJECT_SOURCE_PATH)
|
||||
|
||||
from isaaclab.app import AppLauncher
|
||||
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
|
||||
import cli_args
|
||||
|
||||
parser = argparse.ArgumentParser(description="Evaluate an RSL-RL checkpoint with criterion metrics.")
|
||||
parser.add_argument("--num_envs", type=int, default=250, help="Number of parallel environments.")
|
||||
parser.add_argument("--task", type=str, default="Rough-Deeprobotics-M20-v0", help="Name of the task.")
|
||||
parser.add_argument(
|
||||
"--agent", type=str, default="rsl_rl_cfg_entry_point", help="Name of the RL agent configuration entry point."
|
||||
)
|
||||
parser.add_argument("--seed", type=int, default=None, help="Seed used for the environment.")
|
||||
parser.add_argument(
|
||||
"--eval_type",
|
||||
type=str,
|
||||
default="all",
|
||||
choices=[
|
||||
"all",
|
||||
"upper_limits",
|
||||
"velocity_error",
|
||||
"access_rate",
|
||||
"time_consumption",
|
||||
"energy_consumption",
|
||||
"motion_stability",
|
||||
"stability_margin",
|
||||
"robustness",
|
||||
],
|
||||
)
|
||||
parser.add_argument("--episode_length_s", type=float, default=10.0)
|
||||
parser.add_argument("--command_x", type=float, default=1.0)
|
||||
parser.add_argument("--command_y", type=float, default=0.0)
|
||||
parser.add_argument("--command_yaw", type=float, default=0.0)
|
||||
parser.add_argument("--terrain_size", type=float, default=2.0, help="Evaluation terrain tile size in meters.")
|
||||
parser.add_argument("--success_distance", type=float, default=2.0)
|
||||
parser.add_argument("--start_x", type=float, default=-0.9, help="Evaluation reset x position on the terrain tile.")
|
||||
parser.add_argument("--start_y", type=float, default=0.0, help="Evaluation reset y position on the terrain tile.")
|
||||
parser.add_argument("--num_levels", type=int, default=8, help="Number of difficulty levels for upper_limits.")
|
||||
parser.add_argument("--envs_per_cell", type=int, default=8, help="Number of robots per terrain/difficulty cell.")
|
||||
parser.add_argument(
|
||||
"--terrain_type",
|
||||
type=str,
|
||||
default="all",
|
||||
choices=["all", "stairs_up", "stairs_down", "smooth_slope", "rough_slope", "boxes"],
|
||||
help="Terrain type used by upper_limits.",
|
||||
)
|
||||
parser.add_argument("--tilt_threshold", type=float, default=0.70, help="Tilt threshold in radians for success/stability.")
|
||||
parser.add_argument("--acc_threshold", type=float, default=2.0, help="Base acceleration threshold for instability.")
|
||||
parser.add_argument("--robot_mass", type=float, default=25.482, help="Robot mass used for COT.")
|
||||
parser.add_argument("--output", type=str, default=None, help="Optional output JSON path.")
|
||||
parser.add_argument("--print_json", action="store_true", default=False, help="Print the saved JSON payload to stdout.")
|
||||
parser.add_argument("--verbose", action="store_true", default=False, help="Print per-terrain per-level debug lines.")
|
||||
parser.add_argument(
|
||||
"--include_summary",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="When eval_type=all, also run a second mixed-terrain rollout for summary metrics.",
|
||||
)
|
||||
cli_args.add_rsl_rl_args(parser)
|
||||
AppLauncher.add_app_launcher_args(parser)
|
||||
args_cli, hydra_args = parser.parse_known_args()
|
||||
|
||||
sys.argv = [sys.argv[0]] + hydra_args
|
||||
|
||||
app_launcher = AppLauncher(args_cli)
|
||||
simulation_app = app_launcher.app
|
||||
|
||||
import gymnasium as gym
|
||||
import torch
|
||||
from rsl_rl.runners import OnPolicyRunner
|
||||
|
||||
from isaaclab.envs import (
|
||||
DirectMARLEnv,
|
||||
DirectMARLEnvCfg,
|
||||
DirectRLEnvCfg,
|
||||
ManagerBasedRLEnvCfg,
|
||||
multi_agent_to_single_agent,
|
||||
)
|
||||
from isaaclab.terrains.config.rough import ROUGH_TERRAINS_CFG
|
||||
from isaaclab_rl.rsl_rl import RslRlOnPolicyRunnerCfg, RslRlVecEnvWrapper
|
||||
from isaaclab_tasks.utils import get_checkpoint_path
|
||||
from isaaclab_tasks.utils.hydra import hydra_task_config
|
||||
|
||||
import rl_training
|
||||
import rl_training.tasks # noqa: F401
|
||||
|
||||
EVAL_TYPES = [
|
||||
"upper_limits",
|
||||
"velocity_error",
|
||||
"access_rate",
|
||||
"time_consumption",
|
||||
"energy_consumption",
|
||||
"motion_stability",
|
||||
"stability_margin",
|
||||
"robustness",
|
||||
]
|
||||
|
||||
TERRAIN_LEVELS = {
|
||||
"stairs_up": [0.04, 0.06, 0.08, 0.10, 0.12, 0.14, 0.16, 0.18],
|
||||
"stairs_down": [0.04, 0.06, 0.08, 0.10, 0.12, 0.14, 0.16, 0.18],
|
||||
"smooth_slope": [0.05, 0.10, 0.15, 0.20, 0.25, 0.30, 0.35, 0.40],
|
||||
"rough_slope": [0.05, 0.10, 0.15, 0.20, 0.25, 0.30, 0.35, 0.40],
|
||||
"boxes": [0.04, 0.06, 0.08, 0.10, 0.12, 0.14, 0.16, 0.18],
|
||||
}
|
||||
|
||||
TERRAIN_TO_SUBTERRAIN = {
|
||||
"stairs_up": "pyramid_stairs",
|
||||
"stairs_down": "pyramid_stairs_inv",
|
||||
"smooth_slope": "hf_pyramid_slope",
|
||||
"rough_slope": "random_rough",
|
||||
"boxes": "boxes",
|
||||
}
|
||||
|
||||
|
||||
def apply_eval_terrain_baseline(env_cfg) -> None:
|
||||
"""Use the same evaluation terrain generator for every task."""
|
||||
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)
|
||||
env_cfg.scene.terrain.max_init_terrain_level = None
|
||||
env_cfg.scene.terrain.terrain_generator.size = (args_cli.terrain_size, args_cli.terrain_size)
|
||||
env_cfg.scene.terrain.terrain_generator.border_width = 2.0
|
||||
env_cfg.scene.terrain.terrain_generator.curriculum = False
|
||||
env_cfg.scene.terrain.terrain_generator.num_rows = 5
|
||||
env_cfg.scene.terrain.terrain_generator.num_cols = 5
|
||||
env_cfg.scene.terrain.terrain_generator.use_cache = False
|
||||
for sub_cfg in env_cfg.scene.terrain.terrain_generator.sub_terrains.values():
|
||||
if hasattr(sub_cfg, "platform_width"):
|
||||
sub_cfg.platform_width = min(0.4, args_cli.terrain_size * 0.2)
|
||||
if hasattr(sub_cfg, "border_width"):
|
||||
sub_cfg.border_width = min(0.1, args_cli.terrain_size * 0.05)
|
||||
if hasattr(sub_cfg, "step_width"):
|
||||
sub_cfg.step_width = min(0.25, args_cli.terrain_size * 0.125)
|
||||
if hasattr(sub_cfg, "grid_width"):
|
||||
sub_cfg.grid_width = min(0.3, args_cli.terrain_size * 0.15)
|
||||
|
||||
|
||||
def resolve_checkpoint(log_root_path: str, load_run: str | None, checkpoint: str | None) -> str:
|
||||
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"Unable to find 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"Unable to find run directory: {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 set_fixed_command(env_cfg, command_x: float, command_y: float, command_yaw: float, episode_length_s: float) -> None:
|
||||
command_cfg = env_cfg.commands.base_velocity
|
||||
command_cfg.heading_command = False
|
||||
command_cfg.resampling_time_range = (episode_length_s + 0.01, episode_length_s + 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)
|
||||
|
||||
|
||||
def configure_eval_env(env_cfg, eval_type: str) -> None:
|
||||
env_cfg.scene.num_envs = args_cli.num_envs
|
||||
env_cfg.episode_length_s = args_cli.episode_length_s
|
||||
env_cfg.seed = args_cli.seed if args_cli.seed is not None else 0
|
||||
env_cfg.sim.device = args_cli.device if args_cli.device is not None else env_cfg.sim.device
|
||||
env_cfg.observations.policy.enable_corruption = False
|
||||
set_fixed_command(env_cfg, args_cli.command_x, args_cli.command_y, args_cli.command_yaw, args_cli.episode_length_s)
|
||||
|
||||
apply_eval_terrain_baseline(env_cfg)
|
||||
|
||||
env_cfg.curriculum.terrain_levels = None
|
||||
env_cfg.curriculum.command_levels = None
|
||||
|
||||
for name in (
|
||||
"randomize_rigid_body_material",
|
||||
"randomize_wheel_material",
|
||||
"randomize_support_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",
|
||||
):
|
||||
if hasattr(env_cfg.events, name):
|
||||
setattr(env_cfg.events, name, None)
|
||||
|
||||
if hasattr(env_cfg.events, "randomize_reset_base") and env_cfg.events.randomize_reset_base is not None:
|
||||
env_cfg.events.randomize_reset_base.params = {
|
||||
"pose_range": {
|
||||
"x": (args_cli.start_x, args_cli.start_x),
|
||||
"y": (args_cli.start_y, args_cli.start_y),
|
||||
"z": (0.0, 0.0),
|
||||
"roll": (0.0, 0.0),
|
||||
"pitch": (0.0, 0.0),
|
||||
"yaw": (0.0, 0.0),
|
||||
},
|
||||
"velocity_range": {
|
||||
"x": (0.0, 0.0),
|
||||
"y": (0.0, 0.0),
|
||||
"z": (0.0, 0.0),
|
||||
"roll": (0.0, 0.0),
|
||||
"pitch": (0.0, 0.0),
|
||||
"yaw": (0.0, 0.0),
|
||||
},
|
||||
}
|
||||
|
||||
if hasattr(env_cfg.events, "randomize_push_robot"):
|
||||
env_cfg.events.randomize_push_robot = None
|
||||
if eval_type == "robustness" and hasattr(env_cfg.events, "randomize_push_robot"):
|
||||
# Keep robustness deterministic enough to compare runs, but enable interval pushes.
|
||||
from isaaclab.managers import EventTermCfg as EventTerm
|
||||
import rl_training.tasks.manager_based.locomotion.velocity.mdp as mdp
|
||||
|
||||
env_cfg.events.randomize_push_robot = EventTerm(
|
||||
func=mdp.push_by_setting_velocity,
|
||||
mode="interval",
|
||||
interval_range_s=(2.0, 2.01),
|
||||
params={"velocity_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5)}},
|
||||
)
|
||||
|
||||
|
||||
def configure_upper_limit_env(env_cfg, terrain_types: list[str]) -> None:
|
||||
"""Configure one fixed terrain grid: rows are difficulty levels, columns are terrain types."""
|
||||
configure_eval_env(env_cfg, "all")
|
||||
terrain_generator = env_cfg.scene.terrain.terrain_generator
|
||||
terrain_generator.curriculum = True
|
||||
terrain_generator.num_rows = args_cli.num_levels
|
||||
terrain_generator.num_cols = len(terrain_types)
|
||||
terrain_generator.difficulty_range = (0.0, 1.0)
|
||||
if hasattr(terrain_generator, "seed"):
|
||||
terrain_generator.seed = args_cli.seed if args_cli.seed is not None else 0
|
||||
env_cfg.scene.terrain.max_init_terrain_level = args_cli.num_levels - 1
|
||||
env_cfg.scene.num_envs = args_cli.num_levels * len(terrain_types) * args_cli.envs_per_cell
|
||||
|
||||
selected_subterrains = [TERRAIN_TO_SUBTERRAIN[terrain_type] for terrain_type in terrain_types]
|
||||
original_subterrains = terrain_generator.sub_terrains
|
||||
terrain_generator.sub_terrains = {
|
||||
subterrain_name: copy.deepcopy(original_subterrains[subterrain_name]) for subterrain_name in selected_subterrains
|
||||
}
|
||||
for name, sub_cfg in terrain_generator.sub_terrains.items():
|
||||
sub_cfg.proportion = selected_subterrains.count(name) / len(selected_subterrains)
|
||||
|
||||
for terrain_type in terrain_types:
|
||||
sub_cfg = terrain_generator.sub_terrains[TERRAIN_TO_SUBTERRAIN[terrain_type]]
|
||||
levels = TERRAIN_LEVELS[terrain_type][: args_cli.num_levels]
|
||||
if terrain_type in ("stairs_up", "stairs_down"):
|
||||
sub_cfg.step_height_range = (levels[0], levels[-1])
|
||||
elif terrain_type == "smooth_slope":
|
||||
sub_cfg.slope_range = (levels[0], levels[-1])
|
||||
elif terrain_type == "rough_slope":
|
||||
sub_cfg.noise_range = (levels[0], levels[-1])
|
||||
sub_cfg.noise_step = 0.01
|
||||
elif terrain_type == "boxes":
|
||||
sub_cfg.grid_height_range = (levels[0], levels[-1])
|
||||
|
||||
|
||||
def assign_upper_limit_origins(env, terrain_types: list[str]) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Assign an equal number of robots to every terrain/difficulty cell."""
|
||||
base_env = env.unwrapped
|
||||
terrain = base_env.scene.terrain
|
||||
device = base_env.device
|
||||
num_levels = args_cli.num_levels
|
||||
num_types = len(terrain_types)
|
||||
levels = []
|
||||
types = []
|
||||
for terrain_type_id in range(num_types):
|
||||
for level_id in range(num_levels):
|
||||
levels.extend([level_id] * args_cli.envs_per_cell)
|
||||
types.extend([terrain_type_id] * args_cli.envs_per_cell)
|
||||
levels_tensor = torch.tensor(levels, dtype=torch.long, device=device)
|
||||
types_tensor = torch.tensor(types, dtype=torch.long, device=device)
|
||||
terrain.terrain_levels = levels_tensor
|
||||
terrain.terrain_types = types_tensor
|
||||
terrain.env_origins = terrain.terrain_origins[levels_tensor, types_tensor]
|
||||
env.reset()
|
||||
return levels_tensor, types_tensor
|
||||
|
||||
|
||||
def tilt_from_projected_gravity(projected_gravity_b: torch.Tensor) -> torch.Tensor:
|
||||
z_axis = torch.clamp(-projected_gravity_b[:, 2], min=-1.0, max=1.0)
|
||||
return torch.acos(z_axis)
|
||||
|
||||
|
||||
def support_margin_xy(root_xy: torch.Tensor, foot_xy: torch.Tensor) -> torch.Tensor:
|
||||
min_xy = torch.min(foot_xy, dim=1).values
|
||||
max_xy = torch.max(foot_xy, dim=1).values
|
||||
inside_x = (root_xy[:, 0] >= min_xy[:, 0]) & (root_xy[:, 0] <= max_xy[:, 0])
|
||||
inside_y = (root_xy[:, 1] >= min_xy[:, 1]) & (root_xy[:, 1] <= max_xy[:, 1])
|
||||
inside = inside_x & inside_y
|
||||
inside_dist = torch.minimum(
|
||||
torch.minimum(root_xy[:, 0] - min_xy[:, 0], max_xy[:, 0] - root_xy[:, 0]),
|
||||
torch.minimum(root_xy[:, 1] - min_xy[:, 1], max_xy[:, 1] - root_xy[:, 1]),
|
||||
)
|
||||
outside_dx = torch.clamp(torch.maximum(min_xy[:, 0] - root_xy[:, 0], root_xy[:, 0] - max_xy[:, 0]), min=0.0)
|
||||
outside_dy = torch.clamp(torch.maximum(min_xy[:, 1] - root_xy[:, 1], root_xy[:, 1] - max_xy[:, 1]), min=0.0)
|
||||
outside_dist = torch.sqrt(outside_dx**2 + outside_dy**2)
|
||||
return torch.where(inside, inside_dist, -outside_dist)
|
||||
|
||||
|
||||
def run_rollout(env, policy, include_per_env: bool = False) -> dict[str, object]:
|
||||
base_env = env.unwrapped
|
||||
robot = base_env.scene["robot"]
|
||||
dt = base_env.step_dt
|
||||
device = base_env.device
|
||||
max_steps = int(base_env.max_episode_length)
|
||||
command = torch.tensor([args_cli.command_x, args_cli.command_y, args_cli.command_yaw], device=device)
|
||||
command_xy = command[:2].unsqueeze(0)
|
||||
command_yaw = command[2]
|
||||
command_dir = torch.tensor([args_cli.command_x, args_cli.command_y], device=device)
|
||||
if torch.linalg.norm(command_dir) < 1e-6:
|
||||
command_dir = torch.tensor([1.0, 0.0], device=device)
|
||||
command_dir = command_dir / torch.linalg.norm(command_dir)
|
||||
|
||||
obs = env.get_observations()
|
||||
start_pos = robot.data.root_pos_w[:, :3].clone()
|
||||
last_lin_vel = robot.data.root_lin_vel_w[:, :3].clone()
|
||||
complete_time = torch.full((base_env.num_envs,), float("nan"), device=device)
|
||||
energy = torch.zeros(base_env.num_envs, device=device)
|
||||
route = torch.zeros(base_env.num_envs, device=device)
|
||||
lin_error_sum = torch.zeros(base_env.num_envs, device=device)
|
||||
yaw_error_sum = torch.zeros(base_env.num_envs, device=device)
|
||||
instability_count = torch.zeros(base_env.num_envs, device=device)
|
||||
max_tilt = torch.zeros(base_env.num_envs, device=device)
|
||||
stable_time = torch.zeros(base_env.num_envs, device=device)
|
||||
margin_sum = torch.zeros(base_env.num_envs, device=device)
|
||||
margin_count = 0
|
||||
max_progress = torch.zeros(base_env.num_envs, device=device)
|
||||
reached = torch.zeros(base_env.num_envs, dtype=torch.bool, device=device)
|
||||
ever_fallen = torch.zeros(base_env.num_envs, dtype=torch.bool, device=device)
|
||||
support_steps = torch.zeros(base_env.num_envs, device=device)
|
||||
support_events = torch.zeros(base_env.num_envs, device=device)
|
||||
max_support_time = torch.zeros(base_env.num_envs, device=device)
|
||||
previous_support_active = torch.zeros(base_env.num_envs, dtype=torch.bool, device=device)
|
||||
|
||||
wheel_body_ids, _ = robot.find_bodies(".*_wheel")
|
||||
contact_sensor = base_env.scene.sensors["contact_forces"]
|
||||
support_body_ids, _ = contact_sensor.find_bodies(["base_link", ".*_knee"])
|
||||
|
||||
for step in range(max_steps):
|
||||
pre_progress = torch.sum((robot.data.root_pos_w[:, :2] - start_pos[:, :2]) * command_dir.unsqueeze(0), dim=1)
|
||||
max_progress = torch.maximum(max_progress, pre_progress)
|
||||
reached |= pre_progress >= args_cli.success_distance
|
||||
|
||||
with torch.inference_mode():
|
||||
actions = policy(obs)
|
||||
if isinstance(actions, tuple):
|
||||
actions = actions[0]
|
||||
obs, _, dones, _ = env.step(actions)
|
||||
|
||||
lin_vel_b = robot.data.root_lin_vel_b[:, :2]
|
||||
yaw_vel_b = robot.data.root_ang_vel_b[:, 2]
|
||||
lin_error_sum += torch.linalg.norm(lin_vel_b - command_xy, dim=1)
|
||||
yaw_error_sum += torch.abs(yaw_vel_b - command_yaw)
|
||||
|
||||
joint_power = torch.abs(robot.data.applied_torque * robot.data.joint_vel)
|
||||
energy += torch.sum(joint_power, dim=1) * dt
|
||||
route += torch.linalg.norm(robot.data.root_lin_vel_w[:, :2], dim=1) * dt
|
||||
|
||||
acc = torch.linalg.norm(robot.data.root_lin_vel_w[:, :3] - last_lin_vel, dim=1) / dt
|
||||
instability_count += (acc > args_cli.acc_threshold).float()
|
||||
last_lin_vel = robot.data.root_lin_vel_w[:, :3].clone()
|
||||
|
||||
tilt = tilt_from_projected_gravity(robot.data.projected_gravity_b)
|
||||
max_tilt = torch.maximum(max_tilt, tilt)
|
||||
stable_time += (tilt < args_cli.tilt_threshold).float() * dt
|
||||
|
||||
progress = torch.sum((robot.data.root_pos_w[:, :2] - start_pos[:, :2]) * command_dir.unsqueeze(0), dim=1)
|
||||
max_progress = torch.maximum(max_progress, progress)
|
||||
reached |= progress >= args_cli.success_distance
|
||||
newly_success = torch.isnan(complete_time) & (progress >= args_cli.success_distance)
|
||||
complete_time[newly_success] = (step + 1) * dt
|
||||
ever_fallen |= tilt > args_cli.tilt_threshold
|
||||
|
||||
support_force = torch.linalg.norm(
|
||||
contact_sensor.data.net_forces_w[:, support_body_ids, :], dim=-1
|
||||
)
|
||||
support_time = contact_sensor.data.current_contact_time[:, support_body_ids]
|
||||
support_active = torch.any(support_force > 15.0, dim=1)
|
||||
support_steps += support_active.float()
|
||||
support_events += (support_active & ~previous_support_active).float()
|
||||
max_support_time = torch.maximum(max_support_time, torch.amax(support_time, dim=1))
|
||||
previous_support_active = support_active
|
||||
|
||||
if len(wheel_body_ids) >= 3:
|
||||
foot_xy = robot.data.body_pos_w[:, wheel_body_ids, :2]
|
||||
margin_sum += support_margin_xy(robot.data.root_pos_w[:, :2], foot_xy)
|
||||
margin_count += 1
|
||||
|
||||
tilt = tilt_from_projected_gravity(robot.data.projected_gravity_b)
|
||||
success = reached & ~ever_fallen
|
||||
success_count = torch.sum(success).item()
|
||||
success_den = max(base_env.num_envs, 1)
|
||||
successful_time = complete_time[success & torch.isfinite(complete_time)]
|
||||
avg_time = torch.mean(successful_time).item() if successful_time.numel() > 0 else float("nan")
|
||||
|
||||
cot = energy / ((route + 1e-6) * args_cli.robot_mass * 9.81)
|
||||
metrics = {
|
||||
"num_envs": float(base_env.num_envs),
|
||||
"episode_length_s": float(args_cli.episode_length_s),
|
||||
"terrain_size_m": float(args_cli.terrain_size),
|
||||
"success_distance_m": float(args_cli.success_distance),
|
||||
"success_count": float(success_count),
|
||||
"success_rate": float(success_count / success_den),
|
||||
"avg_final_progress_m": float(torch.mean(max_progress).item()),
|
||||
"avg_time_consumption_s": float(avg_time),
|
||||
"time_efficiency": float(avg_time / (args_cli.success_distance / max(abs(args_cli.command_x), 1e-6)))
|
||||
if successful_time.numel() > 0
|
||||
else float("nan"),
|
||||
"mean_lin_vel_error": float(torch.mean(lin_error_sum / max_steps).item()),
|
||||
"mean_yaw_vel_error": float(torch.mean(yaw_error_sum / max_steps).item()),
|
||||
"avg_energy_j": float(torch.mean(energy).item()),
|
||||
"avg_cost_of_transport": float(torch.mean(cot).item()),
|
||||
"unstability_rate": float(torch.mean(instability_count / max_steps).item()),
|
||||
"avg_stability_margin_m": float(torch.mean(margin_sum / max(margin_count, 1)).item()),
|
||||
"avg_max_tilt_rad": float(torch.mean(max_tilt).item()),
|
||||
"avg_stable_time_s": float(torch.mean(stable_time).item()),
|
||||
"avg_support_contact_ratio": float(torch.mean(support_steps / max_steps).item()),
|
||||
"avg_support_events": float(torch.mean(support_events).item()),
|
||||
"avg_max_support_contact_time_s": float(torch.mean(max_support_time).item()),
|
||||
"successful_runs_with_support_ratio": float(
|
||||
torch.mean((support_events[success] > 0.0).float()).item()
|
||||
)
|
||||
if success_count > 0
|
||||
else float("nan"),
|
||||
}
|
||||
if include_per_env:
|
||||
metrics["_per_env"] = {
|
||||
"success": success.detach().cpu().tolist(),
|
||||
"max_progress": max_progress.detach().cpu().tolist(),
|
||||
"max_tilt": max_tilt.detach().cpu().tolist(),
|
||||
}
|
||||
return metrics
|
||||
|
||||
|
||||
def make_env_and_policy(env_cfg, agent_cfg, resume_path: str):
|
||||
env = gym.make(args_cli.task, cfg=env_cfg)
|
||||
if isinstance(env.unwrapped, DirectMARLEnv):
|
||||
env = multi_agent_to_single_agent(env)
|
||||
env = RslRlVecEnvWrapper(env, clip_actions=agent_cfg.clip_actions)
|
||||
|
||||
runner = OnPolicyRunner(env, agent_cfg.to_dict(), log_dir=None, device=agent_cfg.device)
|
||||
runner.load(resume_path)
|
||||
policy = runner.get_inference_policy(device=env.unwrapped.device)
|
||||
return env, policy
|
||||
|
||||
|
||||
def print_upper_limit_matrix(title: str, terrain_types: list[str], matrix: list[list[float | None]]) -> None:
|
||||
"""Print an upper-limit success matrix with levels as rows and terrain types as columns."""
|
||||
col_width = 14
|
||||
print(f"\n{title}")
|
||||
print("level".ljust(8) + "".join(terrain_type[: col_width - 1].rjust(col_width) for terrain_type in terrain_types))
|
||||
for level_index, row in enumerate(matrix, start=1):
|
||||
cells = []
|
||||
for value in row:
|
||||
cells.append(" -".rjust(col_width) if value is None else f"{value:.2f}".rjust(col_width))
|
||||
print(f"{level_index:<8}" + "".join(cells))
|
||||
|
||||
|
||||
def matrix_to_numpy(matrix: list[list[float | None]]) -> np.ndarray:
|
||||
return np.array([[np.nan if value is None else value for value in row] for row in matrix], dtype=np.float32)
|
||||
|
||||
|
||||
def print_numpy_matrix(title: str, matrix: list[list[float | None]]) -> None:
|
||||
print(f"{title}:")
|
||||
print(np.array2string(matrix_to_numpy(matrix).round(2), precision=2, suppress_small=True))
|
||||
|
||||
|
||||
def print_upper_limit_summary(result: dict[str, object], terrain_types: list[str]) -> None:
|
||||
matrix = result["_matrix"]
|
||||
print_numpy_matrix("success rate", matrix["success_rate"])
|
||||
print_numpy_matrix("success rate accumulated", matrix["success_rate_accumulated"])
|
||||
score = np.array([result[terrain_type]["score"] for terrain_type in terrain_types], dtype=np.float32)
|
||||
print(f"score: {np.array2string(score.round(2), precision=2, suppress_small=True)} of {len(terrain_types)}")
|
||||
|
||||
|
||||
def print_rollout_summary(metrics: dict[str, object]) -> None:
|
||||
success_rate = float(metrics["success_rate"])
|
||||
success_count = int(round(float(metrics["success_count"])))
|
||||
avg_time = float(metrics["avg_time_consumption_s"])
|
||||
avg_time_text = "nan" if np.isnan(avg_time) else f"{avg_time:.2f}"
|
||||
print(f"success rate: {success_rate:.2f}")
|
||||
print(f"success count: {success_count} time consumption: {avg_time_text}")
|
||||
print(f"average power consumption: {float(metrics['avg_energy_j']):.2f}")
|
||||
print(f"average cost of transport: {float(metrics['avg_cost_of_transport']):.2f}")
|
||||
print(f"unstability rate: {float(metrics['unstability_rate']) * 100:.1f}%")
|
||||
print(f"average stability margin: {float(metrics['avg_stability_margin_m']):.2f}")
|
||||
print(
|
||||
f"average max tilt: {float(metrics['avg_max_tilt_rad']):.2f} "
|
||||
f"average max time: {float(metrics['avg_stable_time_s']):.2f}"
|
||||
)
|
||||
support_success_ratio = float(metrics["successful_runs_with_support_ratio"])
|
||||
support_success_text = "nan" if np.isnan(support_success_ratio) else f"{support_success_ratio * 100:.1f}%"
|
||||
print(
|
||||
f"support contact ratio: {float(metrics['avg_support_contact_ratio']) * 100:.1f}% "
|
||||
f"events/episode: {float(metrics['avg_support_events']):.2f} "
|
||||
f"max duration: {float(metrics['avg_max_support_contact_time_s']):.2f}s "
|
||||
f"successful runs using support: {support_success_text}"
|
||||
)
|
||||
|
||||
|
||||
def print_metric_result(eval_type: str, metrics: dict[str, object]) -> None:
|
||||
if eval_type == "velocity_error":
|
||||
print(f"mean lin vel error: {float(metrics['mean_lin_vel_error']):.2f}")
|
||||
print(f"mean yaw vel error: {float(metrics['mean_yaw_vel_error']):.2f}")
|
||||
elif eval_type == "access_rate":
|
||||
print(f"success rate: {float(metrics['success_rate']):.2f}")
|
||||
elif eval_type == "time_consumption":
|
||||
success_count = int(round(float(metrics["success_count"])))
|
||||
avg_time = float(metrics["avg_time_consumption_s"])
|
||||
avg_time_text = "nan" if np.isnan(avg_time) else f"{avg_time:.2f}"
|
||||
print(f"success count: {success_count} time consumption: {avg_time_text}")
|
||||
elif eval_type == "energy_consumption":
|
||||
print(f"average power consumption: {float(metrics['avg_energy_j']):.2f}")
|
||||
print(f"average cost of transport: {float(metrics['avg_cost_of_transport']):.2f}")
|
||||
elif eval_type == "motion_stability":
|
||||
print(f"unstability rate: {float(metrics['unstability_rate']) * 100:.1f}%")
|
||||
elif eval_type == "stability_margin":
|
||||
print(f"average stability margin: {float(metrics['avg_stability_margin_m']):.2f}")
|
||||
elif eval_type == "robustness":
|
||||
print(
|
||||
f"average max tilt: {float(metrics['avg_max_tilt_rad']):.2f} "
|
||||
f"average max time: {float(metrics['avg_stable_time_s']):.2f}"
|
||||
)
|
||||
|
||||
|
||||
def run_upper_limits(env_cfg, agent_cfg, resume_path: str) -> dict[str, object]:
|
||||
terrain_types = list(TERRAIN_LEVELS) if args_cli.terrain_type == "all" else [args_cli.terrain_type]
|
||||
eval_cfg = copy.deepcopy(env_cfg)
|
||||
configure_upper_limit_env(eval_cfg, terrain_types)
|
||||
env, policy = make_env_and_policy(eval_cfg, agent_cfg, resume_path)
|
||||
level_ids, type_ids = assign_upper_limit_origins(env, terrain_types)
|
||||
metrics = run_rollout(env, policy, include_per_env=True)
|
||||
env.close()
|
||||
|
||||
per_env = metrics.pop("_per_env")
|
||||
success = torch.tensor(per_env["success"], dtype=torch.float32)
|
||||
max_progress = torch.tensor(per_env["max_progress"], dtype=torch.float32)
|
||||
max_tilt = torch.tensor(per_env["max_tilt"], dtype=torch.float32)
|
||||
level_ids_cpu = level_ids.cpu()
|
||||
type_ids_cpu = type_ids.cpu()
|
||||
|
||||
result: dict[str, object] = {}
|
||||
for terrain_type_id, terrain_type in enumerate(terrain_types):
|
||||
level_values = TERRAIN_LEVELS[terrain_type][: args_cli.num_levels]
|
||||
rates = []
|
||||
level_metrics = []
|
||||
for level_index, level_value in enumerate(level_values):
|
||||
mask = (type_ids_cpu == terrain_type_id) & (level_ids_cpu == level_index)
|
||||
success_count = float(torch.sum(success[mask]).item())
|
||||
total_count = int(torch.sum(mask).item())
|
||||
success_rate = success_count / max(total_count, 1)
|
||||
avg_progress = float(torch.mean(max_progress[mask]).item()) if total_count > 0 else float("nan")
|
||||
avg_tilt = float(torch.mean(max_tilt[mask]).item()) if total_count > 0 else float("nan")
|
||||
rates.append(success_rate)
|
||||
level_metrics.append(
|
||||
{
|
||||
"level": level_index + 1,
|
||||
"difficulty": level_value,
|
||||
"success_rate": success_rate,
|
||||
"success_count": success_count,
|
||||
"num_envs": total_count,
|
||||
"avg_progress_m": avg_progress,
|
||||
"avg_max_tilt_rad": avg_tilt,
|
||||
}
|
||||
)
|
||||
if args_cli.verbose:
|
||||
print(
|
||||
f"[EVAL] {terrain_type} level={level_index + 1} difficulty={level_value:.3f} "
|
||||
f"success_rate={success_rate:.3f}"
|
||||
)
|
||||
accumulated = []
|
||||
running = 1.0
|
||||
for rate in rates:
|
||||
running *= rate
|
||||
accumulated.append(running)
|
||||
score = 0.0
|
||||
for idx, acc_rate in enumerate(accumulated):
|
||||
next_fail = 1.0 - rates[idx + 1] if idx + 1 < len(rates) else 1.0
|
||||
score += acc_rate * (idx + 1) * next_fail
|
||||
result[terrain_type] = {
|
||||
"levels": level_metrics,
|
||||
"success_rate": rates,
|
||||
"success_rate_accumulated": accumulated,
|
||||
"score": score,
|
||||
"max_level": len(level_values),
|
||||
"benchmark": "fixed_rough_terrain_v1",
|
||||
"terrain_size_m": args_cli.terrain_size,
|
||||
"success_distance_m": args_cli.success_distance,
|
||||
"episode_length_s": args_cli.episode_length_s,
|
||||
}
|
||||
max_levels = max((len(result[terrain_type]["success_rate"]) for terrain_type in terrain_types), default=0)
|
||||
success_matrix: list[list[float | None]] = []
|
||||
accumulated_matrix: list[list[float | None]] = []
|
||||
difficulty_matrix: list[list[float | None]] = []
|
||||
for level_index in range(max_levels):
|
||||
success_row = []
|
||||
accumulated_row = []
|
||||
difficulty_row = []
|
||||
for terrain_type in terrain_types:
|
||||
rates = result[terrain_type]["success_rate"]
|
||||
accumulated = result[terrain_type]["success_rate_accumulated"]
|
||||
levels = result[terrain_type]["levels"]
|
||||
success_row.append(rates[level_index] if level_index < len(rates) else None)
|
||||
accumulated_row.append(accumulated[level_index] if level_index < len(accumulated) else None)
|
||||
difficulty_row.append(levels[level_index]["difficulty"] if level_index < len(levels) else None)
|
||||
success_matrix.append(success_row)
|
||||
accumulated_matrix.append(accumulated_row)
|
||||
difficulty_matrix.append(difficulty_row)
|
||||
|
||||
result["_matrix"] = {
|
||||
"terrain_types": terrain_types,
|
||||
"success_rate": success_matrix,
|
||||
"success_rate_accumulated": accumulated_matrix,
|
||||
"difficulty": difficulty_matrix,
|
||||
"terrain_size_m": args_cli.terrain_size,
|
||||
"success_distance_m": args_cli.success_distance,
|
||||
"episode_length_s": args_cli.episode_length_s,
|
||||
}
|
||||
result["_summary"] = metrics
|
||||
print_upper_limit_summary(result, terrain_types)
|
||||
return result
|
||||
|
||||
|
||||
def filter_metrics(metrics: dict[str, float], eval_type: str) -> dict[str, float]:
|
||||
if eval_type == "all":
|
||||
return metrics
|
||||
keys = {
|
||||
"velocity_error": ("mean_lin_vel_error", "mean_yaw_vel_error"),
|
||||
"access_rate": ("success_count", "success_rate", "avg_final_progress_m"),
|
||||
"time_consumption": ("success_count", "avg_time_consumption_s", "time_efficiency"),
|
||||
"energy_consumption": ("avg_energy_j", "avg_cost_of_transport"),
|
||||
"motion_stability": ("unstability_rate",),
|
||||
"stability_margin": ("avg_stability_margin_m",),
|
||||
"robustness": ("success_rate", "avg_max_tilt_rad", "avg_stable_time_s"),
|
||||
}[eval_type]
|
||||
return {key: metrics[key] for key in keys}
|
||||
|
||||
|
||||
@hydra_task_config(args_cli.task, args_cli.agent)
|
||||
def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agent_cfg: RslRlOnPolicyRunnerCfg):
|
||||
agent_cfg = cli_args.update_rsl_rl_cfg(agent_cfg, args_cli)
|
||||
|
||||
log_root_path = os.path.abspath(os.path.join("logs", "rsl_rl", agent_cfg.experiment_name))
|
||||
resume_path = resolve_checkpoint(log_root_path, agent_cfg.load_run, agent_cfg.load_checkpoint)
|
||||
print(f"[INFO] Evaluating checkpoint: {resume_path}")
|
||||
|
||||
if args_cli.eval_type == "upper_limits":
|
||||
selected = run_upper_limits(env_cfg, agent_cfg, resume_path)
|
||||
elif args_cli.eval_type == "all":
|
||||
upper_limits = run_upper_limits(env_cfg, agent_cfg, resume_path)
|
||||
selected = {"upper_limits": upper_limits}
|
||||
print_rollout_summary(upper_limits["_summary"])
|
||||
if args_cli.include_summary:
|
||||
summary_cfg = copy.deepcopy(env_cfg)
|
||||
configure_eval_env(summary_cfg, "all")
|
||||
env, policy = make_env_and_policy(summary_cfg, agent_cfg, resume_path)
|
||||
summary = run_rollout(env, policy)
|
||||
env.close()
|
||||
selected["summary"] = summary
|
||||
else:
|
||||
eval_type_for_env = "robustness" if args_cli.eval_type == "robustness" else "all"
|
||||
configure_eval_env(env_cfg, eval_type_for_env)
|
||||
env, policy = make_env_and_policy(env_cfg, agent_cfg, resume_path)
|
||||
metrics = run_rollout(env, policy)
|
||||
selected = filter_metrics(metrics, args_cli.eval_type)
|
||||
env.close()
|
||||
print_metric_result(args_cli.eval_type, selected)
|
||||
if args_cli.print_json:
|
||||
print(json.dumps(selected, indent=2, sort_keys=True))
|
||||
|
||||
output = args_cli.output
|
||||
if output is None:
|
||||
eval_dir = os.path.join(os.path.dirname(resume_path), "eval")
|
||||
os.makedirs(eval_dir, exist_ok=True)
|
||||
suffix = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
|
||||
output = os.path.join(eval_dir, f"{args_cli.eval_type}_{suffix}.json")
|
||||
else:
|
||||
output = os.path.abspath(os.path.expanduser(output))
|
||||
output_dir = os.path.dirname(output)
|
||||
if output_dir:
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
with open(output, "w", encoding="utf-8") as f:
|
||||
json.dump({"eval_type": args_cli.eval_type, "checkpoint": resume_path, "metrics": selected}, f, indent=2)
|
||||
print(f"[INFO] Saved evaluation result to: {output}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
simulation_app.close()
|
||||
+23
@@ -0,0 +1,23 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
TASK=${TASK:-Rough-Deeprobotics-M20-v0}
|
||||
CHECKPOINT=${1:-${CHECKPOINT:-}}
|
||||
NUM_ENVS=${NUM_ENVS:-250}
|
||||
DEVICE=${DEVICE:-cuda:0}
|
||||
EVAL_TYPE=${EVAL_TYPE:-all}
|
||||
|
||||
if [[ -z "${CHECKPOINT}" ]]; then
|
||||
echo "Usage: $0 <checkpoint.pt>"
|
||||
echo "Or set CHECKPOINT=/path/to/model.pt"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "Running eval_type=${EVAL_TYPE}"
|
||||
python scripts/reinforcement_learning/rsl_rl/eval.py \
|
||||
--task="${TASK}" \
|
||||
--checkpoint="${CHECKPOINT}" \
|
||||
--num_envs="${NUM_ENVS}" \
|
||||
--device="${DEVICE}" \
|
||||
--eval_type="${EVAL_TYPE}" \
|
||||
--headless
|
||||
@@ -0,0 +1,280 @@
|
||||
# Copyright (c) 2025 Deep Robotics
|
||||
# SPDX-License-Identifier: BSD 3-Clause
|
||||
|
||||
# Copyright (c) 2024-2025 Ziqi Fan
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# Copyright (c) 2024-2025, The Isaac Lab Project Developers.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
"""Script to play a checkpoint if an RL agent from RSL-RL."""
|
||||
|
||||
"""Launch Isaac Sim Simulator first."""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
|
||||
PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", ".."))
|
||||
PROJECT_SOURCE_PATH = os.path.join(PROJECT_ROOT, "source", "rl_training")
|
||||
if PROJECT_SOURCE_PATH not in sys.path:
|
||||
sys.path.insert(0, PROJECT_SOURCE_PATH)
|
||||
|
||||
from isaaclab.app import AppLauncher
|
||||
|
||||
# local imports
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
|
||||
import cli_args
|
||||
from rl_utils import camera_follow
|
||||
|
||||
# add argparse arguments
|
||||
parser = argparse.ArgumentParser(description="Train an RL agent with RSL-RL.")
|
||||
parser.add_argument("--video", action="store_true", default=False, help="Record videos during training.")
|
||||
parser.add_argument("--video_length", type=int, default=200, help="Length of the recorded video (in steps).")
|
||||
parser.add_argument(
|
||||
"--disable_fabric", action="store_true", default=False, help="Disable fabric and use USD I/O operations."
|
||||
)
|
||||
parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.")
|
||||
parser.add_argument("--task", type=str, default=None, help="Name of the task.")
|
||||
parser.add_argument(
|
||||
"--agent", type=str, default="rsl_rl_cfg_entry_point", help="Name of the RL agent configuration entry point."
|
||||
)
|
||||
parser.add_argument("--seed", type=int, default=None, help="Seed used for the environment")
|
||||
parser.add_argument(
|
||||
"--use_pretrained_checkpoint",
|
||||
action="store_true",
|
||||
help="Use the pre-trained checkpoint from Nucleus.",
|
||||
)
|
||||
parser.add_argument("--real-time", action="store_true", default=False, help="Run in real-time, if possible.")
|
||||
parser.add_argument("--keyboard", action="store_true", default=False, help="Whether to use keyboard.")
|
||||
parser.add_argument(
|
||||
"--visual-only",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Display the environment with zero actions without loading an RL checkpoint.",
|
||||
)
|
||||
# append RSL-RL cli arguments
|
||||
cli_args.add_rsl_rl_args(parser)
|
||||
# append AppLauncher cli args
|
||||
AppLauncher.add_app_launcher_args(parser)
|
||||
# parse the arguments
|
||||
args_cli, hydra_args = parser.parse_known_args()
|
||||
# always enable cameras to record video
|
||||
if args_cli.video:
|
||||
args_cli.enable_cameras = True
|
||||
|
||||
# clear out sys.argv for Hydra
|
||||
sys.argv = [sys.argv[0]] + hydra_args
|
||||
|
||||
# launch omniverse app
|
||||
app_launcher = AppLauncher(args_cli)
|
||||
simulation_app = app_launcher.app
|
||||
|
||||
"""Rest everything follows."""
|
||||
|
||||
import gymnasium as gym
|
||||
import time
|
||||
import torch
|
||||
|
||||
from rsl_rl.runners import OnPolicyRunner
|
||||
|
||||
from isaaclab.devices import Se2Keyboard, Se2KeyboardCfg
|
||||
from isaaclab.envs import (
|
||||
DirectMARLEnv,
|
||||
DirectMARLEnvCfg,
|
||||
DirectRLEnvCfg,
|
||||
ManagerBasedRLEnvCfg,
|
||||
multi_agent_to_single_agent,
|
||||
)
|
||||
from isaaclab.managers import ObservationTermCfg as ObsTerm
|
||||
from isaaclab.utils.assets import retrieve_file_path
|
||||
from isaaclab.utils.dict import print_dict
|
||||
from isaaclab.utils.pretrained_checkpoint import get_published_pretrained_checkpoint
|
||||
from isaaclab_rl.rsl_rl import RslRlOnPolicyRunnerCfg, RslRlVecEnvWrapper, export_policy_as_jit, export_policy_as_onnx
|
||||
from isaaclab_tasks.utils import get_checkpoint_path
|
||||
from isaaclab_tasks.utils.hydra import hydra_task_config
|
||||
|
||||
import rl_training
|
||||
import rl_training.tasks # noqa: F401
|
||||
|
||||
print(f"[INFO] Using rl_training package from: {rl_training.__file__}")
|
||||
|
||||
|
||||
@hydra_task_config(args_cli.task, args_cli.agent)
|
||||
def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agent_cfg: RslRlOnPolicyRunnerCfg):
|
||||
"""Play with RSL-RL agent."""
|
||||
task_name = args_cli.task.split(":")[-1]
|
||||
# override configurations with non-hydra CLI arguments
|
||||
agent_cfg = cli_args.update_rsl_rl_cfg(agent_cfg, args_cli)
|
||||
env_cfg.scene.num_envs = args_cli.num_envs if args_cli.num_envs is not None else 50
|
||||
|
||||
# set the environment seed
|
||||
# note: certain randomizations occur in the environment initialization so we set the seed here
|
||||
env_cfg.seed = agent_cfg.seed
|
||||
env_cfg.sim.device = args_cli.device if args_cli.device is not None else env_cfg.sim.device
|
||||
|
||||
# spawn the robot randomly in the grid (instead of their terrain levels)
|
||||
env_cfg.scene.terrain.max_init_terrain_level = None
|
||||
# reduce the number of terrains to save memory
|
||||
if env_cfg.scene.terrain.terrain_generator is not None:
|
||||
env_cfg.scene.terrain.terrain_generator.num_rows = 5
|
||||
env_cfg.scene.terrain.terrain_generator.num_cols = 5
|
||||
env_cfg.scene.terrain.terrain_generator.curriculum = False
|
||||
|
||||
# disable randomization for play
|
||||
env_cfg.observations.policy.enable_corruption = False
|
||||
# Remove perturbations and parameter randomization for reproducible playback.
|
||||
for event_name in (
|
||||
"randomize_rigid_body_material",
|
||||
"randomize_wheel_material",
|
||||
"randomize_support_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",
|
||||
"randomize_push_robot",
|
||||
):
|
||||
if hasattr(env_cfg.events, event_name):
|
||||
setattr(env_cfg.events, event_name, None)
|
||||
env_cfg.curriculum.command_levels = None
|
||||
|
||||
if args_cli.keyboard:
|
||||
env_cfg.scene.num_envs = 1
|
||||
env_cfg.terminations.time_out = None
|
||||
env_cfg.commands.base_velocity.debug_vis = False
|
||||
config = Se2KeyboardCfg(
|
||||
v_x_sensitivity=env_cfg.commands.base_velocity.ranges.lin_vel_x[1]/2,
|
||||
v_y_sensitivity=env_cfg.commands.base_velocity.ranges.lin_vel_y[1],
|
||||
omega_z_sensitivity=env_cfg.commands.base_velocity.ranges.ang_vel_z[1],
|
||||
)
|
||||
controller = Se2Keyboard(config)
|
||||
env_cfg.observations.policy.velocity_commands = ObsTerm(
|
||||
func=lambda env: torch.tensor(controller.advance(), dtype=torch.float32).unsqueeze(0).to(env.device),
|
||||
)
|
||||
|
||||
# specify directory for logging experiments
|
||||
log_root_path = os.path.join("logs", "rsl_rl", agent_cfg.experiment_name)
|
||||
log_root_path = os.path.abspath(log_root_path)
|
||||
if args_cli.visual_only:
|
||||
resume_path = None
|
||||
log_dir = os.path.join(log_root_path, "visual_only")
|
||||
print("[INFO] Visual-only mode: no policy checkpoint will be loaded.")
|
||||
else:
|
||||
print(f"[INFO] Loading experiment from directory: {log_root_path}")
|
||||
if args_cli.use_pretrained_checkpoint:
|
||||
resume_path = get_published_pretrained_checkpoint("rsl_rl", task_name)
|
||||
if not resume_path:
|
||||
print("[INFO] Unfortunately a pre-trained checkpoint is currently unavailable for this task.")
|
||||
return
|
||||
elif args_cli.load_run is not None:
|
||||
resume_path = get_checkpoint_path(log_root_path, agent_cfg.load_run, agent_cfg.load_checkpoint)
|
||||
elif args_cli.checkpoint:
|
||||
resume_path = retrieve_file_path(args_cli.checkpoint)
|
||||
else:
|
||||
resume_path = get_checkpoint_path(log_root_path, agent_cfg.load_run, agent_cfg.load_checkpoint)
|
||||
log_dir = os.path.dirname(resume_path)
|
||||
|
||||
# create isaac environment
|
||||
env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None)
|
||||
|
||||
# convert to single-agent instance if required by the RL algorithm
|
||||
if isinstance(env.unwrapped, DirectMARLEnv):
|
||||
env = multi_agent_to_single_agent(env)
|
||||
|
||||
# wrap for video recording
|
||||
if args_cli.video:
|
||||
video_kwargs = {
|
||||
"video_folder": os.path.join(log_dir, "videos", "play"),
|
||||
"step_trigger": lambda step: step == 0,
|
||||
"video_length": args_cli.video_length,
|
||||
"disable_logger": True,
|
||||
}
|
||||
print("[INFO] Recording videos during training.")
|
||||
print_dict(video_kwargs, nesting=4)
|
||||
env = gym.wrappers.RecordVideo(env, **video_kwargs)
|
||||
|
||||
# wrap around environment for rsl-rl
|
||||
env = RslRlVecEnvWrapper(env, clip_actions=agent_cfg.clip_actions)
|
||||
|
||||
if args_cli.visual_only:
|
||||
policy = None
|
||||
else:
|
||||
print(f"[INFO]: Loading model checkpoint from: {resume_path}")
|
||||
# load previously trained model
|
||||
ppo_runner = OnPolicyRunner(env, agent_cfg.to_dict(), log_dir=None, device=agent_cfg.device)
|
||||
ppo_runner.load(resume_path)
|
||||
|
||||
# obtain the trained policy for inference
|
||||
policy = ppo_runner.get_inference_policy(device=env.unwrapped.device)
|
||||
|
||||
# extract the neural network module
|
||||
# we do this in a try-except to maintain backwards compatibility.
|
||||
try:
|
||||
# version 2.3 onwards
|
||||
policy_nn = ppo_runner.alg.policy
|
||||
except AttributeError:
|
||||
# version 2.2 and below
|
||||
policy_nn = ppo_runner.alg.actor_critic
|
||||
|
||||
# export policy to onnx/jit
|
||||
export_model_dir = os.path.join(os.path.dirname(resume_path), "exported")
|
||||
export_policy_as_onnx(
|
||||
policy=policy_nn,
|
||||
normalizer=None,
|
||||
path=export_model_dir,
|
||||
filename="policy.onnx",
|
||||
)
|
||||
export_policy_as_jit(
|
||||
policy=policy_nn,
|
||||
normalizer=None,
|
||||
path=export_model_dir,
|
||||
filename="policy.pt",
|
||||
)
|
||||
|
||||
dt = env.unwrapped.step_dt
|
||||
# print(dt, "dt")
|
||||
# reset environment
|
||||
|
||||
obs = env.get_observations()
|
||||
|
||||
timestep = 0
|
||||
# simulate environment
|
||||
while simulation_app.is_running():
|
||||
start_time = time.time()
|
||||
# run everything in inference mode
|
||||
with torch.inference_mode():
|
||||
# agent stepping
|
||||
if policy is None:
|
||||
actions = torch.zeros((env.num_envs, env.num_actions), device=env.unwrapped.device)
|
||||
else:
|
||||
actions = policy(obs)
|
||||
|
||||
# env stepping
|
||||
obs, _, _, _ = env.step(actions)
|
||||
if args_cli.video:
|
||||
timestep += 1
|
||||
# Exit the play loop after recording one video
|
||||
if timestep == args_cli.video_length:
|
||||
break
|
||||
|
||||
if args_cli.keyboard:
|
||||
camera_follow(env)
|
||||
|
||||
# time delay for real-time evaluation
|
||||
sleep_time = dt - (time.time() - start_time)
|
||||
if args_cli.real_time and sleep_time > 0:
|
||||
time.sleep(sleep_time)
|
||||
|
||||
# close the simulator
|
||||
env.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# run the main function
|
||||
main()
|
||||
# close sim app
|
||||
simulation_app.close()
|
||||
@@ -0,0 +1,214 @@
|
||||
# Copyright (c) 2025 Deep Robotics
|
||||
# SPDX-License-Identifier: BSD 3-Clause
|
||||
|
||||
# Copyright (c) 2024-2025 Ziqi Fan
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# Copyright (c) 2024-2025, The Isaac Lab Project Developers.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
"""Script to train RL agent with RSL-RL."""
|
||||
|
||||
"""Launch Isaac Sim Simulator first."""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
|
||||
PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", ".."))
|
||||
PROJECT_SOURCE_PATH = os.path.join(PROJECT_ROOT, "source", "rl_training")
|
||||
if PROJECT_SOURCE_PATH not in sys.path:
|
||||
sys.path.insert(0, PROJECT_SOURCE_PATH)
|
||||
|
||||
from isaaclab.app import AppLauncher
|
||||
|
||||
# local imports
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
|
||||
import cli_args
|
||||
|
||||
# add argparse arguments
|
||||
parser = argparse.ArgumentParser(description="Train an RL agent with RSL-RL.")
|
||||
parser.add_argument("--video", action="store_true", default=False, help="Record videos during training.")
|
||||
parser.add_argument("--video_length", type=int, default=200, help="Length of the recorded video (in steps).")
|
||||
parser.add_argument("--video_interval", type=int, default=2000, help="Interval between video recordings (in steps).")
|
||||
parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.")
|
||||
parser.add_argument("--task", type=str, default=None, help="Name of the task.")
|
||||
parser.add_argument(
|
||||
"--agent", type=str, default="rsl_rl_cfg_entry_point", help="Name of the RL agent configuration entry point."
|
||||
)
|
||||
parser.add_argument("--seed", type=int, default=None, help="Seed used for the environment")
|
||||
parser.add_argument("--max_iterations", type=int, default=None, help="RL Policy training iterations.")
|
||||
parser.add_argument(
|
||||
"--distributed", action="store_true", default=False, help="Run training with multiple GPUs or nodes."
|
||||
)
|
||||
# append RSL-RL cli arguments
|
||||
cli_args.add_rsl_rl_args(parser)
|
||||
# append AppLauncher cli args
|
||||
AppLauncher.add_app_launcher_args(parser)
|
||||
args_cli, hydra_args = parser.parse_known_args()
|
||||
|
||||
# always enable cameras to record video
|
||||
if args_cli.video:
|
||||
args_cli.enable_cameras = True
|
||||
|
||||
# clear out sys.argv for Hydra
|
||||
sys.argv = [sys.argv[0]] + hydra_args
|
||||
|
||||
# launch omniverse app
|
||||
app_launcher = AppLauncher(args_cli)
|
||||
simulation_app = app_launcher.app
|
||||
|
||||
"""Rest everything follows."""
|
||||
|
||||
import gymnasium as gym
|
||||
import os
|
||||
import torch
|
||||
from datetime import datetime
|
||||
|
||||
from rsl_rl.runners import OnPolicyRunner
|
||||
|
||||
from isaaclab.envs import (
|
||||
DirectMARLEnv,
|
||||
DirectMARLEnvCfg,
|
||||
DirectRLEnvCfg,
|
||||
ManagerBasedRLEnvCfg,
|
||||
multi_agent_to_single_agent,
|
||||
)
|
||||
from isaaclab.utils.dict import print_dict
|
||||
from isaaclab.utils.io import dump_yaml
|
||||
from isaaclab_rl.rsl_rl import RslRlOnPolicyRunnerCfg, RslRlVecEnvWrapper
|
||||
from isaaclab_tasks.utils import get_checkpoint_path
|
||||
from isaaclab_tasks.utils.hydra import hydra_task_config
|
||||
|
||||
import rl_training
|
||||
import rl_training.tasks # noqa: F401
|
||||
|
||||
print(f"[INFO] Using rl_training package from: {rl_training.__file__}")
|
||||
|
||||
torch.backends.cuda.matmul.allow_tf32 = True
|
||||
torch.backends.cudnn.allow_tf32 = True
|
||||
torch.backends.cudnn.deterministic = False
|
||||
torch.backends.cudnn.benchmark = False
|
||||
|
||||
|
||||
def resolve_resume_path(log_root_path: str, load_run: str | None, load_checkpoint: str | None) -> str:
|
||||
"""Resolve a resume checkpoint, allowing explicit paths across experiment folders."""
|
||||
if load_checkpoint is not None:
|
||||
checkpoint_path = os.path.expanduser(load_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"Unable to find the checkpoint file: {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"Unable to find the run directory: {run_path}")
|
||||
return get_checkpoint_path(os.path.dirname(run_path), os.path.basename(run_path), load_checkpoint)
|
||||
|
||||
return get_checkpoint_path(log_root_path, load_run, load_checkpoint)
|
||||
|
||||
|
||||
@hydra_task_config(args_cli.task, args_cli.agent)
|
||||
def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agent_cfg: RslRlOnPolicyRunnerCfg):
|
||||
"""Train with RSL-RL agent."""
|
||||
# override configurations with non-hydra CLI arguments
|
||||
agent_cfg = cli_args.update_rsl_rl_cfg(agent_cfg, args_cli)
|
||||
env_cfg.scene.num_envs = args_cli.num_envs if args_cli.num_envs is not None else env_cfg.scene.num_envs
|
||||
agent_cfg.max_iterations = (
|
||||
args_cli.max_iterations if args_cli.max_iterations is not None else agent_cfg.max_iterations
|
||||
)
|
||||
|
||||
# set the environment seed
|
||||
# note: certain randomizations occur in the environment initialization so we set the seed here
|
||||
env_cfg.seed = agent_cfg.seed
|
||||
env_cfg.sim.device = args_cli.device if args_cli.device is not None else env_cfg.sim.device
|
||||
|
||||
for reward_name in (
|
||||
"track_diff_wheel_yaw_exp",
|
||||
"lin_vel_xy_yaw_command_l2",
|
||||
"hipx_joint_pos_penalty",
|
||||
"hipy_joint_pos_penalty",
|
||||
):
|
||||
reward_term = getattr(env_cfg.rewards, reward_name, None)
|
||||
reward_weight = None if reward_term is None else reward_term.weight
|
||||
print(f"[INFO] Reward term {reward_name}: weight={reward_weight}")
|
||||
|
||||
# multi-gpu training configuration
|
||||
if args_cli.distributed:
|
||||
env_cfg.sim.device = f"cuda:{app_launcher.local_rank}"
|
||||
agent_cfg.device = f"cuda:{app_launcher.local_rank}"
|
||||
|
||||
# set seed to have diversity in different threads
|
||||
seed = agent_cfg.seed + app_launcher.local_rank
|
||||
env_cfg.seed = seed
|
||||
agent_cfg.seed = seed
|
||||
|
||||
# specify directory for logging experiments
|
||||
log_root_path = os.path.join("logs", "rsl_rl", agent_cfg.experiment_name)
|
||||
log_root_path = os.path.abspath(log_root_path)
|
||||
print(f"[INFO] Logging experiment in directory: {log_root_path}")
|
||||
# specify directory for logging runs: {time-stamp}_{run_name}
|
||||
log_dir = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
|
||||
# The Ray Tune workflow extracts experiment name using the logging line below, hence, do not change it (see PR #2346, comment-2819298849)
|
||||
print(f"Exact experiment name requested from command line: {log_dir}")
|
||||
if agent_cfg.run_name:
|
||||
log_dir += f"_{agent_cfg.run_name}"
|
||||
log_dir = os.path.join(log_root_path, log_dir)
|
||||
|
||||
# create isaac environment
|
||||
env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None)
|
||||
|
||||
# convert to single-agent instance if required by the RL algorithm
|
||||
if isinstance(env.unwrapped, DirectMARLEnv):
|
||||
env = multi_agent_to_single_agent(env)
|
||||
|
||||
# save resume path before creating a new log_dir
|
||||
if agent_cfg.resume or agent_cfg.algorithm.class_name == "Distillation":
|
||||
resume_path = resolve_resume_path(log_root_path, agent_cfg.load_run, agent_cfg.load_checkpoint)
|
||||
|
||||
# wrap for video recording
|
||||
if args_cli.video:
|
||||
video_kwargs = {
|
||||
"video_folder": os.path.join(log_dir, "videos", "train"),
|
||||
"step_trigger": lambda step: step % args_cli.video_interval == 0,
|
||||
"video_length": args_cli.video_length,
|
||||
"disable_logger": True,
|
||||
}
|
||||
print("[INFO] Recording videos during training.")
|
||||
print_dict(video_kwargs, nesting=4)
|
||||
env = gym.wrappers.RecordVideo(env, **video_kwargs)
|
||||
|
||||
# wrap around environment for rsl-rl
|
||||
env = RslRlVecEnvWrapper(env, clip_actions=agent_cfg.clip_actions)
|
||||
|
||||
# create runner from rsl-rl
|
||||
runner = OnPolicyRunner(env, agent_cfg.to_dict(), log_dir=log_dir, device=agent_cfg.device)
|
||||
# write git state to logs
|
||||
runner.add_git_repo_to_log(__file__)
|
||||
# load the checkpoint
|
||||
if agent_cfg.resume or agent_cfg.algorithm.class_name == "Distillation":
|
||||
print(f"[INFO]: Loading model checkpoint from: {resume_path}")
|
||||
# load previously trained model
|
||||
runner.load(resume_path)
|
||||
|
||||
# dump the configuration into log-directory
|
||||
dump_yaml(os.path.join(log_dir, "params", "env.yaml"), env_cfg)
|
||||
dump_yaml(os.path.join(log_dir, "params", "agent.yaml"), agent_cfg)
|
||||
# run training
|
||||
runner.learn(num_learning_iterations=agent_cfg.max_iterations, init_at_random_ep_len=True)
|
||||
|
||||
# close the simulator
|
||||
env.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# run the main function
|
||||
main()
|
||||
# close sim app
|
||||
simulation_app.close()
|
||||
Reference in New Issue
Block a user