Skip to content

Commit 275d0ee

Browse files
louislelaylouislelay
committed
updates rsl rl train script
1 parent de40980 commit 275d0ee

2 files changed

Lines changed: 139 additions & 109 deletions

File tree

‎scripts/rsl_rl/cli_args.py‎

Lines changed: 6 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,4 @@
1-
# Copyright (c) 2024-2025, Muammer Bay (LycheeAI), Louis Le Lay
2-
# All rights reserved.
3-
#
4-
# SPDX-License-Identifier: BSD-3-Clause
5-
#
6-
# Copyright (c) 2022-2025, The Isaac Lab Project Developers.
1+
# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
72
# All rights reserved.
83
#
94
# SPDX-License-Identifier: BSD-3-Clause
@@ -15,7 +10,7 @@
1510
from typing import TYPE_CHECKING
1611

1712
if TYPE_CHECKING:
18-
from isaaclab_rl.rsl_rl import RslRlOnPolicyRunnerCfg
13+
from isaaclab_rl.rsl_rl import RslRlBaseRunnerCfg
1914

2015

2116
def add_rsl_rl_args(parser: argparse.ArgumentParser):
@@ -44,7 +39,7 @@ def add_rsl_rl_args(parser: argparse.ArgumentParser):
4439
)
4540

4641

47-
def parse_rsl_rl_cfg(task_name: str, args_cli: argparse.Namespace) -> RslRlOnPolicyRunnerCfg:
42+
def parse_rsl_rl_cfg(task_name: str, args_cli: argparse.Namespace) -> RslRlBaseRunnerCfg:
4843
"""Parse configuration for RSL-RL agent based on inputs.
4944
5045
Args:
@@ -57,12 +52,12 @@ def parse_rsl_rl_cfg(task_name: str, args_cli: argparse.Namespace) -> RslRlOnPol
5752
from isaaclab_tasks.utils.parse_cfg import load_cfg_from_registry
5853

5954
# load the default configuration
60-
rslrl_cfg: RslRlOnPolicyRunnerCfg = load_cfg_from_registry(task_name, "rsl_rl_cfg_entry_point")
55+
rslrl_cfg: RslRlBaseRunnerCfg = load_cfg_from_registry(task_name, "rsl_rl_cfg_entry_point")
6156
rslrl_cfg = update_rsl_rl_cfg(rslrl_cfg, args_cli)
6257
return rslrl_cfg
6358

6459

65-
def update_rsl_rl_cfg(agent_cfg: RslRlOnPolicyRunnerCfg, args_cli: argparse.Namespace):
60+
def update_rsl_rl_cfg(agent_cfg: RslRlBaseRunnerCfg, args_cli: argparse.Namespace):
6661
"""Update configuration for RSL-RL agent based on inputs.
6762
6863
Args:
@@ -93,4 +88,4 @@ def update_rsl_rl_cfg(agent_cfg: RslRlOnPolicyRunnerCfg, args_cli: argparse.Name
9388
agent_cfg.wandb_project = args_cli.log_project_name
9489
agent_cfg.neptune_project = args_cli.log_project_name
9590

96-
return agent_cfg
91+
return agent_cfg

‎scripts/rsl_rl/play.py‎

Lines changed: 133 additions & 98 deletions
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,14 @@
1-
# Copyright (c) 2024-2025, Muammer Bay (LycheeAI), Louis Le Lay
2-
# All rights reserved.
3-
#
4-
# SPDX-License-Identifier: BSD-3-Clause
5-
#
6-
# Copyright (c) 2022-2025, The Isaac Lab Project Developers.
1+
# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
72
# All rights reserved.
83
#
94
# SPDX-License-Identifier: BSD-3-Clause
105

11-
"""Script to play a checkpoint if an RL agent from RSL-RL."""
6+
"""Script to train RL agent with RSL-RL."""
127

138
"""Launch Isaac Sim Simulator first."""
149

1510
import argparse
11+
import sys
1612

1713
from isaaclab.app import AppLauncher
1814

@@ -23,76 +19,136 @@
2319
parser = argparse.ArgumentParser(description="Train an RL agent with RSL-RL.")
2420
parser.add_argument("--video", action="store_true", default=False, help="Record videos during training.")
2521
parser.add_argument("--video_length", type=int, default=200, help="Length of the recorded video (in steps).")
26-
parser.add_argument(
27-
"--disable_fabric", action="store_true", default=False, help="Disable fabric and use USD I/O operations."
28-
)
22+
parser.add_argument("--video_interval", type=int, default=2000, help="Interval between video recordings (in steps).")
2923
parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.")
3024
parser.add_argument("--task", type=str, default=None, help="Name of the task.")
3125
parser.add_argument(
32-
"--use_pretrained_checkpoint",
33-
action="store_true",
34-
help="Use the pre-trained checkpoint from Nucleus.",
26+
"--agent", type=str, default="rsl_rl_cfg_entry_point", help="Name of the RL agent configuration entry point."
27+
)
28+
parser.add_argument("--seed", type=int, default=None, help="Seed used for the environment")
29+
parser.add_argument("--max_iterations", type=int, default=None, help="RL Policy training iterations.")
30+
parser.add_argument(
31+
"--distributed", action="store_true", default=False, help="Run training with multiple GPUs or nodes."
3532
)
36-
parser.add_argument("--real-time", action="store_true", default=False, help="Run in real-time, if possible.")
33+
parser.add_argument("--export_io_descriptors", action="store_true", default=False, help="Export IO descriptors.")
3734
# append RSL-RL cli arguments
3835
cli_args.add_rsl_rl_args(parser)
3936
# append AppLauncher cli args
4037
AppLauncher.add_app_launcher_args(parser)
41-
args_cli = parser.parse_args()
38+
args_cli, hydra_args = parser.parse_known_args()
39+
4240
# always enable cameras to record video
4341
if args_cli.video:
4442
args_cli.enable_cameras = True
4543

44+
# clear out sys.argv for Hydra
45+
sys.argv = [sys.argv[0]] + hydra_args
46+
4647
# launch omniverse app
4748
app_launcher = AppLauncher(args_cli)
4849
simulation_app = app_launcher.app
4950

51+
"""Check for minimum supported RSL-RL version."""
52+
53+
import importlib.metadata as metadata
54+
import platform
55+
56+
from packaging import version
57+
58+
# check minimum supported rsl-rl version
59+
RSL_RL_VERSION = "3.0.1"
60+
installed_version = metadata.version("rsl-rl-lib")
61+
if version.parse(installed_version) < version.parse(RSL_RL_VERSION):
62+
if platform.system() == "Windows":
63+
cmd = [r".\isaaclab.bat", "-p", "-m", "pip", "install", f"rsl-rl-lib=={RSL_RL_VERSION}"]
64+
else:
65+
cmd = ["./isaaclab.sh", "-p", "-m", "pip", "install", f"rsl-rl-lib=={RSL_RL_VERSION}"]
66+
print(
67+
f"Please install the correct version of RSL-RL.\nExisting version is: '{installed_version}'"
68+
f" and required version is: '{RSL_RL_VERSION}'.\nTo install the correct version, run:"
69+
f"\n\n\t{' '.join(cmd)}\n"
70+
)
71+
exit(1)
72+
5073
"""Rest everything follows."""
5174

75+
import gymnasium as gym
5276
import os
53-
import time
77+
import torch
78+
from datetime import datetime
79+
80+
import omni
81+
from rsl_rl.runners import DistillationRunner, OnPolicyRunner
82+
83+
from isaaclab.envs import (
84+
DirectMARLEnv,
85+
DirectMARLEnvCfg,
86+
DirectRLEnvCfg,
87+
ManagerBasedRLEnvCfg,
88+
multi_agent_to_single_agent,
89+
)
90+
from isaaclab.utils.dict import print_dict
91+
from isaaclab.utils.io import dump_pickle, dump_yaml
92+
93+
from isaaclab_rl.rsl_rl import RslRlBaseRunnerCfg, RslRlVecEnvWrapper
5494

55-
import gymnasium as gym
5695
import isaaclab_tasks # noqa: F401
96+
from isaaclab_tasks.utils import get_checkpoint_path
97+
from isaaclab_tasks.utils.hydra import hydra_task_config
98+
5799
import SO_100.tasks # noqa: F401
58-
import torch
59-
from isaaclab.envs import DirectMARLEnv, multi_agent_to_single_agent
60-
from isaaclab.utils.assets import retrieve_file_path
61-
from isaaclab.utils.dict import print_dict
62-
from isaaclab.utils.pretrained_checkpoint import get_published_pretrained_checkpoint
63-
from isaaclab_rl.rsl_rl import (
64-
RslRlOnPolicyRunnerCfg,
65-
RslRlVecEnvWrapper,
66-
export_policy_as_jit,
67-
export_policy_as_onnx,
68-
)
69-
from isaaclab_tasks.utils import get_checkpoint_path, parse_env_cfg
70-
from rsl_rl.runners import OnPolicyRunner
71100

101+
torch.backends.cuda.matmul.allow_tf32 = True
102+
torch.backends.cudnn.allow_tf32 = True
103+
torch.backends.cudnn.deterministic = False
104+
torch.backends.cudnn.benchmark = False
72105

73-
def main():
74-
"""Play with RSL-RL agent."""
75-
# parse configuration
76-
env_cfg = parse_env_cfg(
77-
args_cli.task, device=args_cli.device, num_envs=args_cli.num_envs, use_fabric=not args_cli.disable_fabric
106+
107+
@hydra_task_config(args_cli.task, args_cli.agent)
108+
def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agent_cfg: RslRlBaseRunnerCfg):
109+
"""Train with RSL-RL agent."""
110+
# override configurations with non-hydra CLI arguments
111+
agent_cfg = cli_args.update_rsl_rl_cfg(agent_cfg, args_cli)
112+
env_cfg.scene.num_envs = args_cli.num_envs if args_cli.num_envs is not None else env_cfg.scene.num_envs
113+
agent_cfg.max_iterations = (
114+
args_cli.max_iterations if args_cli.max_iterations is not None else agent_cfg.max_iterations
78115
)
79-
agent_cfg: RslRlOnPolicyRunnerCfg = cli_args.parse_rsl_rl_cfg(args_cli.task, args_cli)
116+
117+
# set the environment seed
118+
# note: certain randomizations occur in the environment initialization so we set the seed here
119+
env_cfg.seed = agent_cfg.seed
120+
env_cfg.sim.device = args_cli.device if args_cli.device is not None else env_cfg.sim.device
121+
122+
# multi-gpu training configuration
123+
if args_cli.distributed:
124+
env_cfg.sim.device = f"cuda:{app_launcher.local_rank}"
125+
agent_cfg.device = f"cuda:{app_launcher.local_rank}"
126+
127+
# set seed to have diversity in different threads
128+
seed = agent_cfg.seed + app_launcher.local_rank
129+
env_cfg.seed = seed
130+
agent_cfg.seed = seed
80131

81132
# specify directory for logging experiments
82133
log_root_path = os.path.join("logs", "rsl_rl", agent_cfg.experiment_name)
83134
log_root_path = os.path.abspath(log_root_path)
84-
print(f"[INFO] Loading experiment from directory: {log_root_path}")
85-
if args_cli.use_pretrained_checkpoint:
86-
resume_path = get_published_pretrained_checkpoint("rsl_rl", args_cli.task)
87-
if not resume_path:
88-
print("[INFO] Unfortunately a pre-trained checkpoint is currently unavailable for this task.")
89-
return
90-
elif args_cli.checkpoint:
91-
resume_path = retrieve_file_path(args_cli.checkpoint)
135+
print(f"[INFO] Logging experiment in directory: {log_root_path}")
136+
# specify directory for logging runs: {time-stamp}_{run_name}
137+
log_dir = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
138+
# The Ray Tune workflow extracts experiment name using the logging line below, hence, do not change it (see PR #2346, comment-2819298849)
139+
print(f"Exact experiment name requested from command line: {log_dir}")
140+
if agent_cfg.run_name:
141+
log_dir += f"_{agent_cfg.run_name}"
142+
log_dir = os.path.join(log_root_path, log_dir)
143+
144+
# set the IO descriptors output directory if requested
145+
if isinstance(env_cfg, ManagerBasedRLEnvCfg):
146+
env_cfg.export_io_descriptors = args_cli.export_io_descriptors
147+
env_cfg.io_descriptors_output_dir = log_dir
92148
else:
93-
resume_path = get_checkpoint_path(log_root_path, agent_cfg.load_run, agent_cfg.load_checkpoint)
94-
95-
log_dir = os.path.dirname(resume_path)
149+
omni.log.warn(
150+
"IO descriptors are only supported for manager based RL environments. No IO descriptors will be exported."
151+
)
96152

97153
# create isaac environment
98154
env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None)
@@ -101,11 +157,15 @@ def main():
101157
if isinstance(env.unwrapped, DirectMARLEnv):
102158
env = multi_agent_to_single_agent(env)
103159

160+
# save resume path before creating a new log_dir
161+
if agent_cfg.resume or agent_cfg.algorithm.class_name == "Distillation":
162+
resume_path = get_checkpoint_path(log_root_path, agent_cfg.load_run, agent_cfg.load_checkpoint)
163+
104164
# wrap for video recording
105165
if args_cli.video:
106166
video_kwargs = {
107-
"video_folder": os.path.join(log_dir, "videos", "play"),
108-
"step_trigger": lambda step: step == 0,
167+
"video_folder": os.path.join(log_dir, "videos", "train"),
168+
"step_trigger": lambda step: step % args_cli.video_interval == 0,
109169
"video_length": args_cli.video_length,
110170
"disable_logger": True,
111171
}
@@ -116,54 +176,29 @@ def main():
116176
# wrap around environment for rsl-rl
117177
env = RslRlVecEnvWrapper(env, clip_actions=agent_cfg.clip_actions)
118178

119-
print(f"[INFO]: Loading model checkpoint from: {resume_path}")
120-
# load previously trained model
121-
ppo_runner = OnPolicyRunner(env, agent_cfg.to_dict(), log_dir=None, device=agent_cfg.device)
122-
ppo_runner.load(resume_path)
123-
124-
# obtain the trained policy for inference
125-
policy = ppo_runner.get_inference_policy(device=env.unwrapped.device)
126-
127-
# extract the neural network module
128-
# we do this in a try-except to maintain backwards compatibility.
129-
try:
130-
# version 2.3 onwards
131-
policy_nn = ppo_runner.alg.policy
132-
except AttributeError:
133-
# version 2.2 and below
134-
policy_nn = ppo_runner.alg.actor_critic
135-
136-
# export policy to onnx/jit
137-
export_model_dir = os.path.join(os.path.dirname(resume_path), "exported")
138-
export_policy_as_jit(policy_nn, ppo_runner.obs_normalizer, path=export_model_dir, filename="policy.pt")
139-
export_policy_as_onnx(
140-
policy_nn, normalizer=ppo_runner.obs_normalizer, path=export_model_dir, filename="policy.onnx"
141-
)
142-
143-
dt = env.unwrapped.step_dt
144-
145-
# reset environment
146-
obs, _ = env.get_observations()
147-
timestep = 0
148-
# simulate environment
149-
while simulation_app.is_running():
150-
start_time = time.time()
151-
# run everything in inference mode
152-
with torch.inference_mode():
153-
# agent stepping
154-
actions = policy(obs)
155-
# env stepping
156-
obs, _, _, _ = env.step(actions)
157-
if args_cli.video:
158-
timestep += 1
159-
# Exit the play loop after recording one video
160-
if timestep == args_cli.video_length:
161-
break
162-
163-
# time delay for real-time evaluation
164-
sleep_time = dt - (time.time() - start_time)
165-
if args_cli.real_time and sleep_time > 0:
166-
time.sleep(sleep_time)
179+
# create runner from rsl-rl
180+
if agent_cfg.class_name == "OnPolicyRunner":
181+
runner = OnPolicyRunner(env, agent_cfg.to_dict(), log_dir=log_dir, device=agent_cfg.device)
182+
elif agent_cfg.class_name == "DistillationRunner":
183+
runner = DistillationRunner(env, agent_cfg.to_dict(), log_dir=log_dir, device=agent_cfg.device)
184+
else:
185+
raise ValueError(f"Unsupported runner class: {agent_cfg.class_name}")
186+
# write git state to logs
187+
runner.add_git_repo_to_log(__file__)
188+
# load the checkpoint
189+
if agent_cfg.resume or agent_cfg.algorithm.class_name == "Distillation":
190+
print(f"[INFO]: Loading model checkpoint from: {resume_path}")
191+
# load previously trained model
192+
runner.load(resume_path)
193+
194+
# dump the configuration into log-directory
195+
dump_yaml(os.path.join(log_dir, "params", "env.yaml"), env_cfg)
196+
dump_yaml(os.path.join(log_dir, "params", "agent.yaml"), agent_cfg)
197+
dump_pickle(os.path.join(log_dir, "params", "env.pkl"), env_cfg)
198+
dump_pickle(os.path.join(log_dir, "params", "agent.pkl"), agent_cfg)
199+
200+
# run training
201+
runner.learn(num_learning_iterations=agent_cfg.max_iterations, init_at_random_ep_len=True)
167202

168203
# close the simulator
169204
env.close()
@@ -173,4 +208,4 @@ def main():
173208
# run the main function
174209
main()
175210
# close sim app
176-
simulation_app.close()
211+
simulation_app.close()

0 commit comments

Comments
 (0)