228 lines
8.8 KiB
Python
228 lines
8.8 KiB
Python
|
"""Script to play a checkpoint from an RSL-RL agent."""
|
||
|
|
|
||
|
|
"""Launch Isaac Sim Simulator first."""
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
import sys
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
from isaaclab.app import AppLauncher
|
||
|
|
|
||
|
|
# local imports
|
||
|
|
import cli_args # isort: skip
|
||
|
|
|
||
|
|
# ensure repository root is on Python path for Hydra registry imports
|
||
|
|
REPO_ROOT = Path(__file__).resolve().parents[2]
|
||
|
|
if str(REPO_ROOT) not in sys.path:
|
||
|
|
sys.path.append(str(REPO_ROOT))
|
||
|
|
|
||
|
|
# add argparse arguments
|
||
|
|
parser = argparse.ArgumentParser(description="Play an RSL-RL policy checkpoint.")
|
||
|
|
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("--max_steps", type=int, default=None, help="Stop replay after this many policy steps.")
|
||
|
|
|
||
|
|
command_group = parser.add_argument_group("velocity command", description="Velocity command source during replay.")
|
||
|
|
command_group.add_argument(
|
||
|
|
"--command_source",
|
||
|
|
type=str,
|
||
|
|
choices=("random", "fixed", "keyboard"),
|
||
|
|
default="random",
|
||
|
|
help="Use environment-generated, fixed, or keyboard velocity commands.",
|
||
|
|
)
|
||
|
|
command_group.add_argument("--vx", type=float, default=0.3, help="Fixed forward velocity in m/s.")
|
||
|
|
command_group.add_argument("--vy", type=float, default=0.0, help="Fixed lateral velocity in m/s.")
|
||
|
|
command_group.add_argument("--wz", type=float, default=0.0, help="Fixed yaw velocity in rad/s.")
|
||
|
|
command_group.add_argument(
|
||
|
|
"--command_mode",
|
||
|
|
type=str,
|
||
|
|
choices=("step", "ramp"),
|
||
|
|
default="ramp",
|
||
|
|
help="Apply external commands immediately or through acceleration limits.",
|
||
|
|
)
|
||
|
|
command_group.add_argument("--linear_accel", type=float, default=0.8, help="Planar acceleration limit in m/s^2.")
|
||
|
|
command_group.add_argument("--yaw_accel", type=float, default=1.5, help="Yaw acceleration limit in rad/s^2.")
|
||
|
|
command_group.add_argument("--keyboard_vx", type=float, default=0.4, help="Keyboard forward velocity in m/s.")
|
||
|
|
command_group.add_argument("--keyboard_vy", type=float, default=0.2, help="Keyboard lateral velocity in m/s.")
|
||
|
|
command_group.add_argument("--keyboard_wz", type=float, default=0.5, help="Keyboard yaw velocity in rad/s.")
|
||
|
|
|
||
|
|
# 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()
|
||
|
|
if args_cli.linear_accel <= 0.0:
|
||
|
|
parser.error("--linear_accel must be positive.")
|
||
|
|
if args_cli.yaw_accel <= 0.0:
|
||
|
|
parser.error("--yaw_accel must be positive.")
|
||
|
|
if args_cli.max_steps is not None and args_cli.max_steps <= 0:
|
||
|
|
parser.error("--max_steps must be positive.")
|
||
|
|
# always enable cameras to record video
|
||
|
|
# 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 os
|
||
|
|
|
||
|
|
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_rl.rsl_rl import RslRlOnPolicyRunnerCfg, RslRlVecEnvWrapper
|
||
|
|
from isaaclab_tasks.utils import get_checkpoint_path
|
||
|
|
from isaaclab_tasks.utils.hydra import hydra_task_config
|
||
|
|
|
||
|
|
# Import extensions to set up environment tasks
|
||
|
|
import engineai_lab.tasks # noqa: F401
|
||
|
|
|
||
|
|
|
||
|
|
class VelocityCommandController:
|
||
|
|
"""Apply fixed or keyboard commands directly to a velocity command term."""
|
||
|
|
|
||
|
|
def __init__(self, env: RslRlVecEnvWrapper, args: argparse.Namespace):
|
||
|
|
self.env = env
|
||
|
|
self.args = args
|
||
|
|
self.command_term = env.unwrapped.command_manager.get_term("base_velocity")
|
||
|
|
if not hasattr(self.command_term, "vel_command_b"):
|
||
|
|
raise TypeError("The base_velocity command term does not expose a velocity command buffer.")
|
||
|
|
|
||
|
|
self.current_command = torch.zeros(3, device=env.unwrapped.device)
|
||
|
|
self.fixed_command = torch.tensor([args.vx, args.vy, args.wz], device=env.unwrapped.device)
|
||
|
|
self.keyboard = None
|
||
|
|
|
||
|
|
if args.command_source == "keyboard":
|
||
|
|
from isaaclab.devices import Se2Keyboard, Se2KeyboardCfg
|
||
|
|
|
||
|
|
keyboard_cfg = Se2KeyboardCfg(
|
||
|
|
sim_device=str(env.unwrapped.device),
|
||
|
|
v_x_sensitivity=args.keyboard_vx,
|
||
|
|
v_y_sensitivity=args.keyboard_vy,
|
||
|
|
omega_z_sensitivity=args.keyboard_wz,
|
||
|
|
)
|
||
|
|
self.keyboard = Se2Keyboard(keyboard_cfg)
|
||
|
|
print(self.keyboard)
|
||
|
|
else:
|
||
|
|
print(
|
||
|
|
"[INFO] Fixed velocity command: "
|
||
|
|
f"vx={args.vx:.3f} m/s, vy={args.vy:.3f} m/s, wz={args.wz:.3f} rad/s"
|
||
|
|
)
|
||
|
|
print(f"[INFO] External command mode: {args.command_mode}")
|
||
|
|
|
||
|
|
def update(self):
|
||
|
|
target = self.keyboard.advance() if self.keyboard is not None else self.fixed_command
|
||
|
|
if self.args.command_mode == "step":
|
||
|
|
self.current_command.copy_(target)
|
||
|
|
else:
|
||
|
|
self._apply_ramp(target)
|
||
|
|
|
||
|
|
self.command_term.vel_command_b[:] = self.current_command.unsqueeze(0)
|
||
|
|
self.command_term.time_left.fill_(float("inf"))
|
||
|
|
self.command_term.is_standing_env.fill_(False)
|
||
|
|
if hasattr(self.command_term, "is_heading_env"):
|
||
|
|
self.command_term.is_heading_env.fill_(False)
|
||
|
|
|
||
|
|
def _apply_ramp(self, target: torch.Tensor):
|
||
|
|
dt = self.env.unwrapped.step_dt
|
||
|
|
linear_delta = target[:2] - self.current_command[:2]
|
||
|
|
linear_delta_norm = torch.linalg.norm(linear_delta)
|
||
|
|
max_linear_delta = self.args.linear_accel * dt
|
||
|
|
linear_scale = torch.clamp(max_linear_delta / torch.clamp(linear_delta_norm, min=1.0e-6), max=1.0)
|
||
|
|
linear_delta = linear_delta * linear_scale
|
||
|
|
self.current_command[:2] += linear_delta
|
||
|
|
|
||
|
|
yaw_delta = torch.clamp(
|
||
|
|
target[2] - self.current_command[2],
|
||
|
|
min=-self.args.yaw_accel * dt,
|
||
|
|
max=self.args.yaw_accel * dt,
|
||
|
|
)
|
||
|
|
self.current_command[2] += yaw_delta
|
||
|
|
|
||
|
|
|
||
|
|
@hydra_task_config(args_cli.task, "rsl_rl_cfg_entry_point")
|
||
|
|
def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agent_cfg: RslRlOnPolicyRunnerCfg):
|
||
|
|
"""Play with RSL-RL agent."""
|
||
|
|
agent_cfg: RslRlOnPolicyRunnerCfg = cli_args.parse_rsl_rl_cfg(args_cli.task, 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
|
||
|
|
|
||
|
|
# 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}")
|
||
|
|
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=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 around environment for rsl-rl
|
||
|
|
env = RslRlVecEnvWrapper(env)
|
||
|
|
|
||
|
|
command_controller = None
|
||
|
|
if args_cli.command_source != "random":
|
||
|
|
command_controller = VelocityCommandController(env, args_cli)
|
||
|
|
|
||
|
|
# 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)
|
||
|
|
|
||
|
|
# export policy to onnx/jit
|
||
|
|
export_model_dir = os.path.join(os.path.dirname(resume_path), "exported")
|
||
|
|
file_basename = os.path.basename(resume_path).split(".")[0]
|
||
|
|
ppo_runner.export_policy_to_onnx(path=export_model_dir, filename=file_basename+".onnx")
|
||
|
|
|
||
|
|
# reset environment
|
||
|
|
obs = env.get_observations()
|
||
|
|
step_count = 0
|
||
|
|
# simulate environment
|
||
|
|
while simulation_app.is_running() and (args_cli.max_steps is None or step_count < args_cli.max_steps):
|
||
|
|
# run everything in inference mode
|
||
|
|
with torch.inference_mode():
|
||
|
|
if command_controller is not None:
|
||
|
|
command_controller.update()
|
||
|
|
obs = env.get_observations()
|
||
|
|
# agent stepping
|
||
|
|
actions = policy(obs)
|
||
|
|
# env stepping
|
||
|
|
obs, _, _, _ = env.step(actions)
|
||
|
|
step_count += 1
|
||
|
|
|
||
|
|
if args_cli.max_steps is not None:
|
||
|
|
print(f"[INFO] Replay completed after {step_count} policy steps.")
|
||
|
|
|
||
|
|
# close the simulator
|
||
|
|
env.close()
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
# run the main function
|
||
|
|
main()
|
||
|
|
# close sim app
|
||
|
|
simulation_app.close()
|