389 lines
16 KiB
Python
389 lines
16 KiB
Python
|
"""Script to play a trained RSL-RL policy."""
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
import os
|
||
|
|
import sys
|
||
|
|
|
||
|
|
from isaaclab.app import AppLauncher
|
||
|
|
|
||
|
|
parser = argparse.ArgumentParser(description="Play a trained RSL-RL policy.")
|
||
|
|
parser.add_argument("--num_envs", type=int, default=1, 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="Agent configuration entry point.")
|
||
|
|
parser.add_argument("--seed", type=int, default=None, help="Seed used for the environment.")
|
||
|
|
parser.add_argument("--load_run", type=str, default=".*", help="Run folder name or regex to load from.")
|
||
|
|
parser.add_argument(
|
||
|
|
"--checkpoint",
|
||
|
|
type=str,
|
||
|
|
default="model_.*.pt",
|
||
|
|
help="Checkpoint filename regex under logs, or a direct path to a .pt file.",
|
||
|
|
)
|
||
|
|
parser.add_argument("--num_steps", type=int, default=500, help="Number of inference steps to simulate.")
|
||
|
|
parser.add_argument(
|
||
|
|
"--disable_resets",
|
||
|
|
action="store_true",
|
||
|
|
default=False,
|
||
|
|
help="Disable environment reset during visual play so short or unstable policies do not instantly jump back to the start pose.",
|
||
|
|
)
|
||
|
|
parser.add_argument("--export_io_descriptors", action="store_true", default=False, help="Export IO descriptors.")
|
||
|
|
parser.add_argument("--export_onnx", action="store_true", default=False, help="Export policy.onnx from the loaded checkpoint.")
|
||
|
|
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 inspect
|
||
|
|
import os
|
||
|
|
import torch
|
||
|
|
from rsl_rl.runners import OnPolicyRunner
|
||
|
|
try:
|
||
|
|
from rsl_rl.algorithms import PPO
|
||
|
|
except ImportError:
|
||
|
|
PPO = None
|
||
|
|
try:
|
||
|
|
from tensordict import TensorDict
|
||
|
|
except ImportError:
|
||
|
|
TensorDict = None
|
||
|
|
|
||
|
|
from isaaclab.envs import DirectMARLEnv, DirectMARLEnvCfg, DirectRLEnvCfg, ManagerBasedRLEnvCfg, multi_agent_to_single_agent
|
||
|
|
from isaaclab_rl.rsl_rl import RslRlBaseRunnerCfg
|
||
|
|
|
||
|
|
from isaaclab_tasks.utils import get_checkpoint_path
|
||
|
|
from isaaclab_tasks.utils.hydra import hydra_task_config
|
||
|
|
|
||
|
|
import stackforce_simready_2foot_lab.tasks # noqa: F401
|
||
|
|
|
||
|
|
import os
|
||
|
|
import torch
|
||
|
|
|
||
|
|
|
||
|
|
class _StackForceOnnxPolicy(torch.nn.Module):
|
||
|
|
def __init__(self, policy):
|
||
|
|
super().__init__()
|
||
|
|
self.policy = policy
|
||
|
|
|
||
|
|
def forward(self, obs):
|
||
|
|
actions = self.policy(obs)
|
||
|
|
if isinstance(actions, (tuple, list)):
|
||
|
|
return actions[0]
|
||
|
|
if isinstance(actions, dict):
|
||
|
|
if "actions" in actions:
|
||
|
|
return actions["actions"]
|
||
|
|
if "action" in actions:
|
||
|
|
return actions["action"]
|
||
|
|
return next(iter(actions.values()))
|
||
|
|
return actions
|
||
|
|
|
||
|
|
|
||
|
|
def _stackforce_policy_obs_tensor(obs):
|
||
|
|
if isinstance(obs, dict):
|
||
|
|
obs = obs["policy"] if "policy" in obs else next(iter(obs.values()))
|
||
|
|
elif not isinstance(obs, torch.Tensor) and hasattr(obs, "get"):
|
||
|
|
try:
|
||
|
|
candidate = obs.get("policy")
|
||
|
|
except Exception:
|
||
|
|
candidate = None
|
||
|
|
if candidate is not None:
|
||
|
|
obs = candidate
|
||
|
|
if not isinstance(obs, torch.Tensor):
|
||
|
|
raise TypeError(f"ONNX export requires a tensor policy observation, got {type(obs)!r}")
|
||
|
|
if hasattr(obs, "detach"):
|
||
|
|
obs = obs.detach()
|
||
|
|
if obs.dim() == 1:
|
||
|
|
obs = obs.unsqueeze(0)
|
||
|
|
elif obs.shape[0] > 1:
|
||
|
|
obs = obs[:1]
|
||
|
|
return obs.contiguous()
|
||
|
|
|
||
|
|
|
||
|
|
def stackforce_export_policy_as_onnx(policy, obs, output_dir, file_name="policy.onnx", opset=17):
|
||
|
|
os.makedirs(output_dir, exist_ok=True)
|
||
|
|
output_path = os.path.join(output_dir, file_name)
|
||
|
|
sample_obs = _stackforce_policy_obs_tensor(obs)
|
||
|
|
module = _StackForceOnnxPolicy(policy).to(sample_obs.device).eval()
|
||
|
|
with torch.no_grad():
|
||
|
|
torch.onnx.export(
|
||
|
|
module,
|
||
|
|
sample_obs,
|
||
|
|
output_path,
|
||
|
|
input_names=["obs"],
|
||
|
|
output_names=["actions"],
|
||
|
|
dynamic_axes={"obs": {0: "batch"}, "actions": {0: "batch"}},
|
||
|
|
opset_version=opset,
|
||
|
|
)
|
||
|
|
print(f"Exported ONNX policy to: {output_path}")
|
||
|
|
return output_path
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
def _runner_uses_obs_groups():
|
||
|
|
try:
|
||
|
|
source = inspect.getsource(OnPolicyRunner)
|
||
|
|
if PPO is not None and hasattr(PPO, "construct_algorithm"):
|
||
|
|
source += "\n" + inspect.getsource(PPO.construct_algorithm)
|
||
|
|
except OSError:
|
||
|
|
return False
|
||
|
|
return "resolve_obs_groups" in source or '"obs_groups"' in source or "'obs_groups'" in source
|
||
|
|
|
||
|
|
|
||
|
|
def _runner_uses_split_actor_critic():
|
||
|
|
try:
|
||
|
|
source = inspect.getsource(OnPolicyRunner)
|
||
|
|
if PPO is not None and hasattr(PPO, "construct_algorithm"):
|
||
|
|
source += "\n" + inspect.getsource(PPO.construct_algorithm)
|
||
|
|
except OSError:
|
||
|
|
return False
|
||
|
|
return 'cfg["actor"]' in source or "cfg['actor']" in source or 'cfg["critic"]' in source or "cfg['critic']" in source
|
||
|
|
|
||
|
|
|
||
|
|
def _runner_expects_privileged_step():
|
||
|
|
try:
|
||
|
|
source = inspect.getsource(OnPolicyRunner.learn)
|
||
|
|
except OSError:
|
||
|
|
return True
|
||
|
|
return "privileged_obs" in source or "critic_obs" in source
|
||
|
|
|
||
|
|
|
||
|
|
def _format_rsl_rl_obs(obs_dict, use_obs_groups):
|
||
|
|
if not use_obs_groups:
|
||
|
|
return obs_dict["policy"]
|
||
|
|
if TensorDict is not None and not isinstance(obs_dict, TensorDict):
|
||
|
|
first_obs = next(iter(obs_dict.values()))
|
||
|
|
return TensorDict(dict(obs_dict), batch_size=[first_obs.shape[0]], device=first_obs.device)
|
||
|
|
return obs_dict
|
||
|
|
|
||
|
|
|
||
|
|
class LegacyRslRlVecEnvWrapper:
|
||
|
|
def __init__(self, env, clip_actions=None):
|
||
|
|
self.env = env
|
||
|
|
self.clip_actions = clip_actions
|
||
|
|
self.use_obs_groups = _runner_uses_obs_groups()
|
||
|
|
self.return_privileged_obs = _runner_expects_privileged_step()
|
||
|
|
self.num_envs = env.unwrapped.num_envs
|
||
|
|
self.device = env.unwrapped.device
|
||
|
|
self.max_episode_length = env.unwrapped.max_episode_length
|
||
|
|
self.cfg = env.unwrapped.cfg
|
||
|
|
self.num_actions = gym.spaces.flatdim(env.unwrapped.single_action_space)
|
||
|
|
obs_dict, extras = self.env.reset()
|
||
|
|
self.obs_buf = _format_rsl_rl_obs(obs_dict, self.use_obs_groups)
|
||
|
|
self.privileged_obs_buf = obs_dict.get("critic")
|
||
|
|
self.num_obs = obs_dict["policy"].shape[-1]
|
||
|
|
self.num_privileged_obs = self.privileged_obs_buf.shape[-1] if self.privileged_obs_buf is not None else None
|
||
|
|
self.rew_buf = torch.zeros(self.num_envs, device=self.device)
|
||
|
|
self.reset_buf = torch.zeros(self.num_envs, dtype=torch.long, device=self.device)
|
||
|
|
self.episode_length_buf = env.unwrapped.episode_length_buf
|
||
|
|
self.extras = extras
|
||
|
|
|
||
|
|
def get_observations(self):
|
||
|
|
return self.obs_buf
|
||
|
|
|
||
|
|
def get_privileged_observations(self):
|
||
|
|
return self.privileged_obs_buf
|
||
|
|
|
||
|
|
def reset(self, env_ids=None):
|
||
|
|
del env_ids
|
||
|
|
obs_dict, extras = self.env.reset()
|
||
|
|
self.obs_buf = _format_rsl_rl_obs(obs_dict, self.use_obs_groups)
|
||
|
|
self.privileged_obs_buf = obs_dict.get("critic")
|
||
|
|
self.extras = extras
|
||
|
|
if not self.return_privileged_obs:
|
||
|
|
return self.obs_buf
|
||
|
|
return self.obs_buf, self.privileged_obs_buf
|
||
|
|
|
||
|
|
def step(self, actions):
|
||
|
|
if self.clip_actions is not None:
|
||
|
|
actions = torch.clamp(actions, -self.clip_actions, self.clip_actions)
|
||
|
|
obs_dict, rewards, terminated, truncated, extras = self.env.step(actions)
|
||
|
|
dones = (terminated | truncated).to(dtype=torch.long)
|
||
|
|
if not self.env.unwrapped.cfg.is_finite_horizon:
|
||
|
|
extras["time_outs"] = truncated
|
||
|
|
episode_log = dict(extras.get("episode", {}))
|
||
|
|
if torch.any(dones.bool()) and "log" in extras:
|
||
|
|
episode_log.update(extras["log"])
|
||
|
|
episode_log["Step_Reward/mean"] = torch.mean(rewards.detach())
|
||
|
|
episode_log["Step_Reward/abs_mean"] = torch.mean(torch.abs(rewards.detach()))
|
||
|
|
extras["episode"] = episode_log
|
||
|
|
self.obs_buf = _format_rsl_rl_obs(obs_dict, self.use_obs_groups)
|
||
|
|
self.privileged_obs_buf = obs_dict.get("critic")
|
||
|
|
self.rew_buf = rewards
|
||
|
|
self.reset_buf = dones
|
||
|
|
self.extras = extras
|
||
|
|
if not self.return_privileged_obs:
|
||
|
|
return self.obs_buf, rewards, dones, extras
|
||
|
|
return self.obs_buf, self.privileged_obs_buf, rewards, dones, extras
|
||
|
|
|
||
|
|
def close(self):
|
||
|
|
return self.env.close()
|
||
|
|
|
||
|
|
|
||
|
|
def _runner_expects_nested_runner():
|
||
|
|
try:
|
||
|
|
source = inspect.getsource(OnPolicyRunner)
|
||
|
|
except OSError:
|
||
|
|
return True
|
||
|
|
return 'train_cfg["runner"]' in source or "train_cfg['runner']" in source
|
||
|
|
|
||
|
|
|
||
|
|
def _runner_uses_nested_class_name():
|
||
|
|
try:
|
||
|
|
source = inspect.getsource(OnPolicyRunner)
|
||
|
|
except OSError:
|
||
|
|
return False
|
||
|
|
return (
|
||
|
|
'algorithm"]["class_name' in source
|
||
|
|
or "algorithm']['class_name" in source
|
||
|
|
or 'policy_cfg.pop("class_name")' in source
|
||
|
|
or "policy_cfg.pop('class_name')" in source
|
||
|
|
or 'self.policy_cfg.pop("class_name")' in source
|
||
|
|
or "self.policy_cfg.pop('class_name')" in source
|
||
|
|
or "resolve_callable" in source
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def _runner_uses_split_actor_critic():
|
||
|
|
try:
|
||
|
|
source = inspect.getsource(OnPolicyRunner)
|
||
|
|
if PPO is not None and hasattr(PPO, "construct_algorithm"):
|
||
|
|
source += "\n" + inspect.getsource(PPO.construct_algorithm)
|
||
|
|
except OSError:
|
||
|
|
return False
|
||
|
|
return 'cfg["actor"]' in source or "cfg['actor']" in source or 'cfg["critic"]' in source or "cfg['critic']" in source
|
||
|
|
|
||
|
|
|
||
|
|
def to_compatible_rsl_rl_cfg(agent_cfg):
|
||
|
|
data = agent_cfg.to_dict() if hasattr(agent_cfg, "to_dict") else dict(agent_cfg)
|
||
|
|
allowed_policy_keys = {
|
||
|
|
"actor_hidden_dims",
|
||
|
|
"critic_hidden_dims",
|
||
|
|
"activation",
|
||
|
|
"init_noise_std",
|
||
|
|
"clip_actions",
|
||
|
|
"actor_obs_normalization",
|
||
|
|
"critic_obs_normalization",
|
||
|
|
}
|
||
|
|
allowed_algorithm_keys = {
|
||
|
|
"num_learning_epochs",
|
||
|
|
"num_mini_batches",
|
||
|
|
"clip_param",
|
||
|
|
"gamma",
|
||
|
|
"lam",
|
||
|
|
"value_loss_coef",
|
||
|
|
"entropy_coef",
|
||
|
|
"learning_rate",
|
||
|
|
"max_grad_norm",
|
||
|
|
"use_clipped_value_loss",
|
||
|
|
"schedule",
|
||
|
|
"desired_kl",
|
||
|
|
"use_spo",
|
||
|
|
}
|
||
|
|
if "runner" in data and "policy" in data and "algorithm" in data:
|
||
|
|
runner_cfg = dict(data["runner"])
|
||
|
|
policy_cfg = {key: value for key, value in dict(data["policy"]).items() if key in allowed_policy_keys or key == "class_name"}
|
||
|
|
algorithm_cfg = {key: value for key, value in dict(data["algorithm"]).items() if key in allowed_algorithm_keys or key == "class_name"}
|
||
|
|
else:
|
||
|
|
policy_cfg = {key: value for key, value in dict(data["policy"]).items() if key in allowed_policy_keys}
|
||
|
|
algorithm_cfg = {key: value for key, value in dict(data["algorithm"]).items() if key in allowed_algorithm_keys}
|
||
|
|
runner_cfg = {key: value for key, value in data.items() if key not in {"policy", "algorithm", "class_name"}}
|
||
|
|
|
||
|
|
runner_cfg.setdefault("num_steps_per_env", getattr(agent_cfg, "num_steps_per_env", 24))
|
||
|
|
runner_cfg.setdefault("max_iterations", getattr(agent_cfg, "max_iterations", 1500))
|
||
|
|
runner_cfg.setdefault("save_interval", getattr(agent_cfg, "save_interval", 50))
|
||
|
|
runner_cfg.setdefault("obs_groups", {"policy": ["policy"], "critic": ["policy"]})
|
||
|
|
runner_cfg.setdefault("experiment_name", getattr(agent_cfg, "experiment_name", "stackforce"))
|
||
|
|
runner_cfg.setdefault("run_name", getattr(agent_cfg, "run_name", ""))
|
||
|
|
runner_cfg.setdefault("resume", getattr(agent_cfg, "resume", False))
|
||
|
|
runner_cfg.setdefault("load_run", getattr(agent_cfg, "load_run", ".*"))
|
||
|
|
runner_cfg.setdefault("checkpoint", getattr(agent_cfg, "load_checkpoint", "model_.*.pt"))
|
||
|
|
|
||
|
|
if _runner_uses_nested_class_name():
|
||
|
|
policy_cfg.setdefault("class_name", "ActorCritic")
|
||
|
|
algorithm_cfg.setdefault("class_name", "PPO")
|
||
|
|
else:
|
||
|
|
runner_cfg.setdefault("policy_class_name", "ActorCritic")
|
||
|
|
runner_cfg.setdefault("algorithm_class_name", "PPO")
|
||
|
|
|
||
|
|
if _runner_expects_nested_runner():
|
||
|
|
return {"runner": runner_cfg, "policy": policy_cfg, "algorithm": algorithm_cfg}
|
||
|
|
if _runner_uses_split_actor_critic():
|
||
|
|
algorithm_cfg.setdefault("class_name", "PPO")
|
||
|
|
algorithm_cfg.pop("use_spo", None)
|
||
|
|
actor_cfg = {
|
||
|
|
"class_name": "MLPModel",
|
||
|
|
"hidden_dims": policy_cfg.get("actor_hidden_dims", [256, 256, 128]),
|
||
|
|
"activation": policy_cfg.get("activation", "elu"),
|
||
|
|
"obs_normalization": policy_cfg.get("actor_obs_normalization", False),
|
||
|
|
"distribution_cfg": {
|
||
|
|
"class_name": "GaussianDistribution",
|
||
|
|
"init_std": policy_cfg.get("init_noise_std", 1.0),
|
||
|
|
},
|
||
|
|
}
|
||
|
|
critic_cfg = {
|
||
|
|
"class_name": "MLPModel",
|
||
|
|
"hidden_dims": policy_cfg.get("critic_hidden_dims", [256, 256, 128]),
|
||
|
|
"activation": policy_cfg.get("activation", "elu"),
|
||
|
|
"obs_normalization": policy_cfg.get("critic_obs_normalization", False),
|
||
|
|
}
|
||
|
|
runner_cfg.pop("policy_class_name", None)
|
||
|
|
runner_cfg.pop("algorithm_class_name", None)
|
||
|
|
runner_cfg["obs_groups"] = {"actor": ["policy"], "critic": ["policy"], "policy": ["policy"]}
|
||
|
|
runner_cfg.setdefault("multi_gpu", None)
|
||
|
|
return {**runner_cfg, "actor": actor_cfg, "critic": critic_cfg, "algorithm": algorithm_cfg}
|
||
|
|
return {**runner_cfg, "policy": policy_cfg, "algorithm": algorithm_cfg}
|
||
|
|
|
||
|
|
|
||
|
|
@hydra_task_config(args_cli.task, args_cli.agent)
|
||
|
|
def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agent_cfg: RslRlBaseRunnerCfg):
|
||
|
|
env_cfg.scene.num_envs = args_cli.num_envs if args_cli.num_envs is not None else env_cfg.scene.num_envs
|
||
|
|
env_cfg.seed = args_cli.seed if args_cli.seed is not None else agent_cfg.seed
|
||
|
|
env_cfg.sim.device = args_cli.device if args_cli.device is not None else env_cfg.sim.device
|
||
|
|
if isinstance(env_cfg, ManagerBasedRLEnvCfg):
|
||
|
|
env_cfg.export_io_descriptors = args_cli.export_io_descriptors
|
||
|
|
elif args_cli.disable_resets and hasattr(env_cfg, "visual_disable_resets"):
|
||
|
|
env_cfg.visual_disable_resets = True
|
||
|
|
|
||
|
|
log_root_path = os.path.abspath(os.path.join("logs", "rsl_rl", agent_cfg.experiment_name))
|
||
|
|
checkpoint_arg = args_cli.checkpoint
|
||
|
|
if os.path.isfile(checkpoint_arg):
|
||
|
|
resume_path = os.path.abspath(checkpoint_arg)
|
||
|
|
else:
|
||
|
|
resume_path = get_checkpoint_path(log_root_path, args_cli.load_run, checkpoint_arg)
|
||
|
|
print(f"[INFO]: Loading model checkpoint from: {resume_path}")
|
||
|
|
|
||
|
|
env = gym.make(args_cli.task, cfg=env_cfg)
|
||
|
|
if isinstance(env.unwrapped, DirectMARLEnv):
|
||
|
|
env = multi_agent_to_single_agent(env)
|
||
|
|
|
||
|
|
legacy_agent_cfg = to_compatible_rsl_rl_cfg(agent_cfg)
|
||
|
|
wrapped_env = LegacyRslRlVecEnvWrapper(env, clip_actions=getattr(agent_cfg, "clip_actions", None))
|
||
|
|
runner = OnPolicyRunner(wrapped_env, legacy_agent_cfg, log_dir=None, device=env.unwrapped.device)
|
||
|
|
runner.load(resume_path, load_optimizer=False)
|
||
|
|
policy = runner.get_inference_policy(device=env.unwrapped.device)
|
||
|
|
|
||
|
|
obs = wrapped_env.get_observations()
|
||
|
|
if args_cli.export_onnx:
|
||
|
|
stackforce_export_policy_as_onnx(
|
||
|
|
policy,
|
||
|
|
obs,
|
||
|
|
os.path.join(os.path.dirname(resume_path), "exported", "policies"),
|
||
|
|
)
|
||
|
|
steps = 0
|
||
|
|
with torch.inference_mode():
|
||
|
|
while simulation_app.is_running():
|
||
|
|
actions = policy(obs)
|
||
|
|
step_result = wrapped_env.step(actions)
|
||
|
|
obs = step_result[0]
|
||
|
|
steps += 1
|
||
|
|
if args_cli.num_steps > 0 and steps >= args_cli.num_steps:
|
||
|
|
break
|
||
|
|
|
||
|
|
wrapped_env.close()
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|
||
|
|
simulation_app.close()
|