cyy/go2_rl_robotlab
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

init commit.

Browse Source
This commit is contained in:
wertyuilife
2026-03-23 03:17:59 +08:00
commit fce9166e87
99 changed files with 13778 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

96
scripts/reinforcement_learning/rsl_rl/cli_args.py Normal file
View File

@@ -0,0 +1,96 @@
# 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 RslRlBaseRunnerCfg
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) -> RslRlBaseRunnerCfg:
"""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: RslRlBaseRunnerCfg = 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: RslRlBaseRunnerCfg, 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.experiment_name is not None:
agent_cfg.experiment_name = args_cli.experiment_name
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

228
scripts/reinforcement_learning/rsl_rl/play.py Normal file
View File

@@ -0,0 +1,228 @@
# 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
from isaaclab.app import AppLauncher
# local imports
import cli_args # isort: skip
from utils import export_cts_policy_as_jit, export_cts_policy_as_onnx
# 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.")
# 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 scripts.reinforcement_learning.utils import camera_follow
from rsl_rl.runners import DistillationRunner, OnPolicyRunner, OnPolicyRunnerCTS
from isaaclab.devices import Se2Keyboard, Se2KeyboardCfg
from isaaclab.envs import (
DirectMARLEnv,
DirectMARLEnvCfg,
DirectRLEnvCfg,
ManagerBasedRLEnvCfg,
multi_agent_to_single_agent,
)
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 RslRlBaseRunnerCfg, 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 robot_lab.tasks # noqa: F401
@hydra_task_config(args_cli.task, args_cli.agent)
def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agent_cfg: RslRlBaseRunnerCfg):
"""Play with RSL-RL agent."""
# grab task name for checkpoint path
task_name = args_cli.task.split(":")[-1]
# override configurations with non-hydra CLI arguments
agent_cfg: RslRlBaseRunnerCfg = 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 64
# 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
# disable randomization for play
env_cfg.observations.policy.enable_corruption = False
# remove random pushing
env_cfg.events.randomize_apply_external_force_torque = None
env_cfg.events.push_robot = None
env_cfg.curriculum.command_levels_lin_vel = None
env_cfg.curriculum.command_levels_ang_vel = None
# 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] 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
raise NotImplementedError("Pre-trained checkpoint retrieval is disabled temporarily.")
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)
# set the log directory for the environment (works for all environment types)
env_cfg.log_dir = 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)
# 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)
print(f"[INFO]: Loading model checkpoint from: {resume_path}")
# load previously trained model
if agent_cfg.class_name == "OnPolicyRunner":
runner = OnPolicyRunner(env, agent_cfg.to_dict(), log_dir=None, device=agent_cfg.device)
elif agent_cfg.class_name == "DistillationRunner":
runner = DistillationRunner(env, agent_cfg.to_dict(), log_dir=None, device=agent_cfg.device)
elif agent_cfg.class_name == "OnPolicyRunnerCTS":
runner = OnPolicyRunnerCTS(env, agent_cfg.to_dict(), log_dir=log_dir, device=agent_cfg.device)
else:
raise ValueError(f"Unsupported runner class: {agent_cfg.class_name}")
runner.load(resume_path)
# obtain the trained policy for inference
policy = 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 = runner.alg.policy
except AttributeError:
# version 2.2 and below
policy_nn = runner.alg.actor_critic
# extract the normalizer
if hasattr(policy_nn, "actor_obs_normalizer"):
normalizer = policy_nn.actor_obs_normalizer
elif hasattr(policy_nn, "student_obs_normalizer"):
normalizer = policy_nn.student_obs_normalizer
else:
normalizer = None
# export policy to onnx/jit
export_model_dir = os.path.join(os.path.dirname(resume_path), "exported")
if agent_cfg.class_name == "OnPolicyRunnerCTS":
export_cts_policy_as_jit(policy_nn, actor_obs_normalizer=policy_nn.actor_obs_normalizer, single_obs_normalizer=policy_nn.single_obs_normalizer, path=export_model_dir, filename="policy.pt")
export_cts_policy_as_onnx(policy_nn, actor_obs_normalizer=policy_nn.actor_obs_normalizer, single_obs_normalizer=policy_nn.single_obs_normalizer, path=export_model_dir, filename="policy.onnx")
else:
export_policy_as_jit(policy_nn, normalizer=normalizer, path=export_model_dir, filename="policy.pt")
export_policy_as_onnx(policy_nn, normalizer=normalizer, path=export_model_dir, filename="policy.onnx")
dt = env.unwrapped.step_dt
# env.unwrapped.eye = (1.1, 3.3, 0.9)
# 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
actions = policy(obs)
# env stepping
obs, _, dones, _ = env.step(actions)
# reset recurrent states for episodes that have terminated
policy_nn.reset(dones)
if args_cli.video:
timestep += 1
# Exit the play loop after recording one video
if timestep == args_cli.video_length:
break
# 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()

218
scripts/reinforcement_learning/rsl_rl/train.py Normal file
View File

@@ -0,0 +1,218 @@
# 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
from isaaclab.app import AppLauncher
# local imports
import cli_args # isort: skip
# 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."
)
parser.add_argument("--export_io_descriptors", action="store_true", default=False, help="Export IO descriptors.")
# 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
"""Check for minimum supported RSL-RL version."""
import importlib.metadata as metadata
from packaging import version
# check minimum supported rsl-rl version
RSL_RL_VERSION = "3.0.1"
installed_version = metadata.version("rsl-rl-lib")
if version.parse(installed_version) < version.parse(RSL_RL_VERSION):
cmd = [r"python", "-m", "pip", "install", f"rsl-rl-lib=={RSL_RL_VERSION}"]
print(
f"Please install the correct version of RSL-RL.\nExisting version is: '{installed_version}'"
f" and required version is: '{RSL_RL_VERSION}'.\nTo install the correct version, run:"
f"\n\n\t{' '.join(cmd)}\n"
)
exit(1)
"""Rest everything follows."""
import gymnasium as gym
import torch
from datetime import datetime
# local imports
from utils import Logger
import omni
from rsl_rl.runners import DistillationRunner, OnPolicyRunner, OnPolicyRunnerCTS
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 RslRlBaseRunnerCfg, RslRlVecEnvWrapper
from isaaclab_tasks.utils import get_checkpoint_path
from isaaclab_tasks.utils.hydra import hydra_task_config
import robot_lab.tasks # noqa: F401
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
torch.backends.cudnn.deterministic = False
torch.backends.cudnn.benchmark = False
@hydra_task_config(args_cli.task, args_cli.agent)
def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agent_cfg: RslRlBaseRunnerCfg):
"""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
# check for invalid combination of CPU device with distributed training
if args_cli.distributed and args_cli.device is not None and "cpu" in args_cli.device:
raise ValueError(
"Distributed training is not supported when using CPU device. "
"Please use GPU device (e.g., --device cuda) for distributed training."
)
# 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)
# set the IO descriptors export flag if requested
if isinstance(env_cfg, ManagerBasedRLEnvCfg):
env_cfg.export_io_descriptors = args_cli.export_io_descriptors
else:
omni.log.warn(
"IO descriptors are only supported for manager based RL environments. No IO descriptors will be exported."
)
# set the log directory for the environment (works for all environment types)
env_cfg.log_dir = 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 = get_checkpoint_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
if agent_cfg.class_name == "OnPolicyRunner":
runner = OnPolicyRunner(env, agent_cfg.to_dict(), log_dir=log_dir, device=agent_cfg.device)
elif agent_cfg.class_name == "DistillationRunner":
runner = DistillationRunner(env, agent_cfg.to_dict(), log_dir=log_dir, device=agent_cfg.device)
elif agent_cfg.class_name == "OnPolicyRunnerCTS":
runner = OnPolicyRunnerCTS(env, agent_cfg.to_dict(), log_dir=log_dir, device=agent_cfg.device)
else:
raise ValueError(f"Unsupported runner class: {agent_cfg.class_name}")
# 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)
sys.stdout = Logger(os.path.join(log_dir, "train.log"))
# 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()

177
scripts/reinforcement_learning/rsl_rl/utils.py Normal file
View File

@@ -0,0 +1,177 @@
# base version: IsaacLab/source/isaaclab_rl/isaaclab_rl/rsl_rl/exporter.py
import copy
import os
import torch
import re
import os
import sys
from typing import NamedTuple
"Script to log terminal output to a file, stripping ANSI escape codes."
class Logger:
def __init__(self, filename):
self.terminal = sys.stdout
os.makedirs(os.path.dirname(filename), exist_ok=True)
self.log = open(filename, 'w', encoding='utf-8')
self.ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')
def write(self, message):
clean_message = self.ansi_escape.sub('', message)
self.terminal.write(message)
self.log.write(clean_message)
self.log.flush()
def flush(self):
self.terminal.flush()
self.log.flush()
# Inputs of CTS Policy is a TensorDict with 'policy' and 'single_obs' keys, we simulate this with a NamedTuple.
class CTSPolicyInputs(NamedTuple):
policy: torch.Tensor
single_obs: torch.Tensor
def export_cts_policy_as_jit(policy: object, actor_obs_normalizer: object | None, single_obs_normalizer: object | None, path: str, filename="policy.pt"):
"""Export CTS policy into a Torch JIT file.
Args:
policy: The CTS policy torch module.
actor_obs_normalizer: The empirical normalizer module for actor observations. If None, Identity is used.
single_obs_normalizer: The empirical normalizer module for single observations. If None, Identity is used.
path: The path to the saving directory.
filename: The name of exported JIT file. Defaults to "policy.pt".
"""
policy_exporter = _TorchPolicyExporter(policy, actor_obs_normalizer, single_obs_normalizer)
policy_exporter.export(path, filename)
def export_cts_policy_as_onnx(
policy: object, path: str, actor_obs_normalizer: object | None = None, single_obs_normalizer: object | None = None, filename="policy.onnx", verbose=False
):
"""Export CTS policy into a Torch ONNX file.
Args:
policy: The CTS policy torch module.
actor_obs_normalizer: The empirical normalizer module for actor observations. If None, Identity is used.
single_obs_normalizer: The empirical normalizer module for single observations. If None, Identity is used.
path: The path to the saving directory.
filename: The name of exported ONNX file. Defaults to "policy.onnx".
verbose: Whether to print the model summary. Defaults to False.
"""
if not os.path.exists(path):
os.makedirs(path, exist_ok=True)
policy_exporter = _OnnxPolicyExporter(policy, actor_obs_normalizer, single_obs_normalizer, verbose)
policy_exporter.export(path, filename)
"""
Helper Classes - Private.
"""
class _TorchPolicyExporter(torch.nn.Module):
"""Exporter of actor-critic into JIT file."""
def __init__(self, policy, actor_obs_normalizer=None, single_obs_normalizer=None):
assert not policy.is_recurrent, "CTS policy should not be recurrent"
super().__init__()
# copy policy parameters
if hasattr(policy, "actor"):
self.actor = copy.deepcopy(policy.actor)
elif hasattr(policy, "student"):
self.actor = copy.deepcopy(policy.student)
else:
raise ValueError("Policy does not have an actor/student module.")
self.student_moe_encoder = copy.deepcopy(policy.student_moe_encoder)
self.state_dependent_std = policy.state_dependent_std
# copy normalizer if exists
if actor_obs_normalizer:
self.actor_obs_normalizer = copy.deepcopy(actor_obs_normalizer)
else:
self.actor_obs_normalizer = torch.nn.Identity()
if single_obs_normalizer:
self.single_obs_normalizer = copy.deepcopy(single_obs_normalizer)
else:
self.single_obs_normalizer = torch.nn.Identity()
def forward(self, x: CTSPolicyInputs):
single_obs = self.single_obs_normalizer(x.single_obs)
obs_a = self.actor_obs_normalizer(x.policy)
latent, _ = self.student_moe_encoder(obs_a)
latent_and_obs = torch.cat([latent, single_obs], dim=-1)
if self.state_dependent_std:
return self.actor(latent_and_obs)[..., 0, :]
else:
return self.actor(latent_and_obs)
@torch.jit.export
def reset(self):
pass
def export(self, path, filename):
os.makedirs(path, exist_ok=True)
path = os.path.join(path, filename)
self.to("cpu")
traced_script_module = torch.jit.script(self)
traced_script_module.save(path)
class _OnnxPolicyExporter(torch.nn.Module):
"""Exporter of actor-critic into ONNX file."""
def __init__(self, policy, actor_obs_normalizer=None, single_obs_normalizer=None, verbose=False):
assert not policy.is_recurrent, "CTS policy should not be recurrent"
super().__init__()
self.verbose = verbose
# copy policy parameters
if hasattr(policy, "actor"):
self.actor = copy.deepcopy(policy.actor)
elif hasattr(policy, "student"):
self.actor = copy.deepcopy(policy.student)
else:
raise ValueError("Policy does not have an actor/student module.")
self.student_moe_encoder = copy.deepcopy(policy.student_moe_encoder)
self.num_single_obs = policy.num_single_obs
self.num_actor_obs = policy.num_actor_obs
self.state_dependent_std = policy.state_dependent_std
# copy normalizer if exists
if actor_obs_normalizer:
self.actor_obs_normalizer = copy.deepcopy(actor_obs_normalizer)
else:
self.actor_obs_normalizer = torch.nn.Identity()
if single_obs_normalizer:
self.single_obs_normalizer = copy.deepcopy(single_obs_normalizer)
else:
self.single_obs_normalizer = torch.nn.Identity()
def forward(self, history, single_obs):
single_obs = self.single_obs_normalizer(single_obs)
obs_a = self.actor_obs_normalizer(history)
latent, _ = self.student_moe_encoder(obs_a)
latent_and_obs = torch.cat([latent, single_obs], dim=-1)
if self.state_dependent_std:
return self.actor(latent_and_obs)[..., 0, :]
else:
return self.actor(latent_and_obs)
def export(self, path, filename):
self.to("cpu")
self.eval()
opset_version = 18 # was 11, but it caused problems with linux-aarch, and 18 worked well across all systems.
torch.onnx.export(
self,
(torch.zeros(1, self.num_actor_obs), torch.zeros(1, self.num_single_obs)),
os.path.join(path, filename),
export_params=True,
opset_version=opset_version,
verbose=self.verbose,
input_names=["obs"],
output_names=["actions"],
dynamic_axes={},
)

62
scripts/tools/clean_trash.py Normal file
View File

@@ -0,0 +1,62 @@
# Copyright (c) 2024-2025 Ziqi Fan
# SPDX-License-Identifier: Apache-2.0
import os
import re
import shutil
def clean_trash(folder_path):
"""
Delete folders that meet the following conditions:
1. Contain `events.out.*` files.
2. Also meet:
- No `.pt` files, or
- Less than 3 `.pt` files.
:param folder_path: Target folder path
"""
if not os.path.exists(folder_path):
print(f"Folder {folder_path} does not exist.")
return
# List of folders to delete
folders_to_delete = []
# Traverse subfolders
for root, dirs, files in os.walk(folder_path):
# Regex match `events.out.*` and `.pt` files
event_pattern = re.compile(r"events\.out.*")
model_pattern = re.compile(r"model_\d+\.pt")
event_files = [f for f in files if event_pattern.match(f)]
model_files = [f for f in files if model_pattern.match(f)]
# Check delete conditions: have `events.out.*` files, but no `.pt` files or less than 3 `.pt` files
if event_files and (len(model_files) < 3):
folders_to_delete.append(os.path.abspath(root))
# If there are folders that meet the conditions, prompt and delete
if folders_to_delete:
print("The following folders contain `events.out.*` and meet the deletion conditions, they will be deleted:")
for folder in folders_to_delete:
print(f" - {folder}")
# Confirm deletion
confirm = input("Confirm deletion of these folders? (y/n): ").strip().lower()
if confirm == "y":
for folder in folders_to_delete:
shutil.rmtree(folder) # Delete the entire folder
print(f"Deleted: {folder}")
print("All folders that meet the conditions have been deleted.")
else:
print("Deletion operation canceled.")
else:
print("No folders meet the conditions, no need to delete.")
# Example call
if __name__ == "__main__":
# folder = input("Please enter the target folder path: ").strip()
folder = "logs"
clean_trash(folder)

142
scripts/tools/convert_mjcf.py Normal file
View File

@@ -0,0 +1,142 @@
# Copyright (c) 2024-2025 Ziqi Fan
# SPDX-License-Identifier: Apache-2.0
# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause
"""
Utility to convert a MJCF into USD format.
MuJoCo XML Format (MJCF) is an XML file format used in MuJoCo to describe all elements of a robot.
For more information, see: http://www.mujoco.org/book/XMLreference.html
This script uses the MJCF importer extension from Isaac Sim (``isaacsim.asset.importer.mjcf``) to convert
a MJCF asset into USD format. It is designed as a convenience script for command-line use. For more information
on the MJCF importer, see the documentation for the extension:
https://docs.isaacsim.omniverse.nvidia.com/latest/robot_setup/ext_isaacsim_asset_importer_mjcf.html
positional arguments:
input The path to the input URDF file.
output The path to store the USD file.
optional arguments:
-h, --help Show this help message and exit
--fix-base Fix the base to where it is imported. (default: False)
--import-sites Import sites by parse <site> tag. (default: True)
--make-instanceable Make the asset instanceable for efficient cloning. (default: False)
"""
"""Launch Isaac Sim Simulator first."""
import argparse
from isaaclab.app import AppLauncher
# add argparse arguments
parser = argparse.ArgumentParser(description="Utility to convert a MJCF into USD format.")
parser.add_argument("input", type=str, help="The path to the input MJCF file.")
parser.add_argument("output", type=str, help="The path to store the USD file.")
parser.add_argument("--fix-base", action="store_true", default=False, help="Fix the base to where it is imported.")
parser.add_argument(
"--import-sites", action="store_true", default=False, help="Import sites by parsing the <site> tag."
)
parser.add_argument(
"--make-instanceable",
action="store_true",
default=False,
help="Make the asset instanceable for efficient cloning.",
)
# append AppLauncher cli args
AppLauncher.add_app_launcher_args(parser)
# parse the arguments
args_cli = parser.parse_args()
# launch omniverse app
app_launcher = AppLauncher(args_cli)
simulation_app = app_launcher.app
"""Rest everything follows."""
import contextlib
import os
import carb
import isaacsim.core.utils.stage as stage_utils
import omni.kit.app
from isaaclab.sim.converters import MjcfConverter, MjcfConverterCfg
from isaaclab.utils.assets import check_file_path
from isaaclab.utils.dict import print_dict
def main():
# check valid file path
mjcf_path = args_cli.input
if not os.path.isabs(mjcf_path):
mjcf_path = os.path.abspath(mjcf_path)
if not check_file_path(mjcf_path):
raise ValueError(f"Invalid file path: {mjcf_path}")
# create destination path
dest_path = args_cli.output
if not os.path.isabs(dest_path):
dest_path = os.path.abspath(dest_path)
# create the converter configuration
mjcf_converter_cfg = MjcfConverterCfg(
asset_path=mjcf_path,
usd_dir=os.path.dirname(dest_path),
usd_file_name=os.path.basename(dest_path),
fix_base=args_cli.fix_base,
import_sites=args_cli.import_sites,
force_usd_conversion=True,
make_instanceable=args_cli.make_instanceable,
)
# Print info
print("-" * 80)
print("-" * 80)
print(f"Input MJCF file: {mjcf_path}")
print("MJCF importer config:")
print_dict(mjcf_converter_cfg.to_dict(), nesting=0)
print("-" * 80)
print("-" * 80)
# Create mjcf converter and import the file
mjcf_converter = MjcfConverter(mjcf_converter_cfg)
# print output
print("MJCF importer output:")
print(f"Generated USD file: {mjcf_converter.usd_path}")
print("-" * 80)
print("-" * 80)
# Determine if there is a GUI to update:
# acquire settings interface
carb_settings_iface = carb.settings.get_settings()
# read flag for whether a local GUI is enabled
local_gui = carb_settings_iface.get("/app/window/enabled")
# read flag for whether livestreaming GUI is enabled
livestream_gui = carb_settings_iface.get("/app/livestream/enabled")
# Simulate scene (if not headless)
if local_gui or livestream_gui:
# Open the stage with USD
stage_utils.open_stage(mjcf_converter.usd_path)
# Reinitialize the simulation
app = omni.kit.app.get_app_interface()
# Run simulation
with contextlib.suppress(KeyboardInterrupt):
while app.is_running():
# perform step
app.update()
if __name__ == "__main__":
# run the main function
main()
# close sim app
simulation_app.close()

166
scripts/tools/convert_urdf.py Normal file
View File

@@ -0,0 +1,166 @@
# Copyright (c) 2024-2025 Ziqi Fan
# SPDX-License-Identifier: Apache-2.0
# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause
"""
Utility to convert a URDF into USD format.
Unified Robot Description Format (URDF) is an XML file format used in ROS to describe all elements of
a robot. For more information, see: http://wiki.ros.org/urdf
This script uses the URDF importer extension from Isaac Sim (``isaacsim.asset.importer.urdf``) to convert a
URDF asset into USD format. It is designed as a convenience script for command-line use. For more
information on the URDF importer, see the documentation for the extension:
https://docs.isaacsim.omniverse.nvidia.com/latest/robot_setup/ext_isaacsim_asset_importer_urdf.html
positional arguments:
input The path to the input URDF file.
output The path to store the USD file.
optional arguments:
-h, --help Show this help message and exit
--merge-joints Consolidate links that are connected by fixed joints. (default: False)
--fix-base Fix the base to where it is imported. (default: False)
--joint-stiffness The stiffness of the joint drive. (default: 100.0)
--joint-damping The damping of the joint drive. (default: 1.0)
--joint-target-type The type of control to use for the joint drive. (default: "position")
"""
"""Launch Isaac Sim Simulator first."""
import argparse
from isaaclab.app import AppLauncher
# add argparse arguments
parser = argparse.ArgumentParser(description="Utility to convert a URDF into USD format.")
parser.add_argument("input", type=str, help="The path to the input URDF file.")
parser.add_argument("output", type=str, help="The path to store the USD file.")
parser.add_argument(
"--merge-joints",
action="store_true",
default=False,
help="Consolidate links that are connected by fixed joints.",
)
parser.add_argument("--fix-base", action="store_true", default=False, help="Fix the base to where it is imported.")
parser.add_argument(
"--joint-stiffness",
type=float,
default=100.0,
help="The stiffness of the joint drive.",
)
parser.add_argument(
"--joint-damping",
type=float,
default=1.0,
help="The damping of the joint drive.",
)
parser.add_argument(
"--joint-target-type",
type=str,
default="position",
choices=["position", "velocity", "none"],
help="The type of control to use for the joint drive.",
)
# append AppLauncher cli args
AppLauncher.add_app_launcher_args(parser)
# parse the arguments
args_cli = parser.parse_args()
# launch omniverse app
app_launcher = AppLauncher(args_cli)
simulation_app = app_launcher.app
"""Rest everything follows."""
import contextlib
import os
import carb
import isaacsim.core.utils.stage as stage_utils
import omni.kit.app
from isaaclab.sim.converters import UrdfConverter, UrdfConverterCfg
from isaaclab.utils.assets import check_file_path
from isaaclab.utils.dict import print_dict
def main():
# check valid file path
urdf_path = args_cli.input
if not os.path.isabs(urdf_path):
urdf_path = os.path.abspath(urdf_path)
if not check_file_path(urdf_path):
raise ValueError(f"Invalid file path: {urdf_path}")
# create destination path
dest_path = args_cli.output
if not os.path.isabs(dest_path):
dest_path = os.path.abspath(dest_path)
# Create Urdf converter config
urdf_converter_cfg = UrdfConverterCfg(
asset_path=urdf_path,
usd_dir=os.path.dirname(dest_path),
usd_file_name=os.path.basename(dest_path),
fix_base=args_cli.fix_base,
merge_fixed_joints=args_cli.merge_joints,
force_usd_conversion=True,
joint_drive=UrdfConverterCfg.JointDriveCfg(
gains=UrdfConverterCfg.JointDriveCfg.PDGainsCfg(
stiffness=args_cli.joint_stiffness,
damping=args_cli.joint_damping,
),
target_type=args_cli.joint_target_type,
),
)
# Print info
print("-" * 80)
print("-" * 80)
print(f"Input URDF file: {urdf_path}")
print("URDF importer config:")
print_dict(urdf_converter_cfg.to_dict(), nesting=0)
print("-" * 80)
print("-" * 80)
# Create Urdf converter and import the file
urdf_converter = UrdfConverter(urdf_converter_cfg)
# print output
print("URDF importer output:")
print(f"Generated USD file: {urdf_converter.usd_path}")
print("-" * 80)
print("-" * 80)
# Determine if there is a GUI to update:
# acquire settings interface
carb_settings_iface = carb.settings.get_settings()
# read flag for whether a local GUI is enabled
local_gui = carb_settings_iface.get("/app/window/enabled")
# read flag for whether livestreaming GUI is enabled
livestream_gui = carb_settings_iface.get("/app/livestream/enabled")
# Simulate scene (if not headless)
if local_gui or livestream_gui:
# Open the stage with USD
stage_utils.open_stage(urdf_converter.usd_path)
# Reinitialize the simulation
app = omni.kit.app.get_app_interface()
# Run simulation
with contextlib.suppress(KeyboardInterrupt):
while app.is_running():
# perform step
app.update()
if __name__ == "__main__":
# run the main function
main()
# close sim app
simulation_app.close()

77
scripts/tools/list_envs.py Normal file
View File

@@ -0,0 +1,77 @@
# 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 print all the available environments in Isaac Lab.
The script iterates over all registered environments and stores the details in a table.
It prints the name of the environment, the entry point and the config file.
All the environments are registered in the `isaaclab_tasks` extension. They start
with `Isaac` in their name.
"""
"""Launch Isaac Sim Simulator first."""
from isaaclab.app import AppLauncher
# launch omniverse app
app_launcher = AppLauncher(headless=True)
simulation_app = app_launcher.app
"""Rest everything follows."""
import gymnasium as gym
import textwrap
from prettytable import PrettyTable
import robot_lab.tasks # noqa: F401
def main():
"""Print all environments registered in `isaaclab_tasks` extension."""
# print all the available environments
table = PrettyTable(["S. No.", "Task Name", "Entry Point", "Config"])
table.title = "Available Environments in Isaac Lab"
# set alignment of table columns
table.align["Task Name"] = "l"
table.align["Entry Point"] = "l"
table.align["Config"] = "l"
table.hrules = 1
# set max width for text wrapping
max_width = 50
# count of environments
index = 0
# acquire all Isaac environments names
for task_spec in gym.registry.values():
if "RobotLab" in task_spec.id:
# wrap long text in each column before adding it to the table
task_name = textwrap.fill(task_spec.id, max_width)
entry_point = textwrap.fill(task_spec.entry_point, max_width)
config = textwrap.fill(task_spec.kwargs["env_cfg_entry_point"], max_width)
# add details to table
table.add_row([index + 1, task_name, entry_point, config])
# increment count
index += 1
print(table)
if __name__ == "__main__":
try:
# run the main function
main()
except Exception as e:
raise e
finally:
# close the app
simulation_app.close()