1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
| """
使用 RSL-RL v1.0.2 推理/预览 Berkeley Humanoid Lite 的脚本。
结构严格仿照 train.py，适配 Isaac Sim 5.1。
"""
import argparse
import os
import sys
import torch
from isaaclab.app import AppLauncher
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
import cli_args
parser = argparse.ArgumentParser(description="Play/Inference with RSL-RL agent.")
parser.add_argument("--num_envs", type=int, default=1, help="仿真环境数量(预览通常为1)")
parser.add_argument("--task", type=str, default=None, help="任务名称")
parser.add_argument("--seed", type=int, default=None, help="随机种子")
cli_args.add_rsl_rl_args(parser)
AppLauncher.add_app_launcher_args(parser)
args_cli, hydra_args = parser.parse_known_args()
args_cli.headless = False
sys.argv = [sys.argv[0]] + hydra_args
app_launcher = AppLauncher(args_cli)
simulation_app = app_launcher.app
import gymnasium as gym
from omegaconf import OmegaConf
from rsl_rl.runners import OnPolicyRunner
from rsl_rl.algorithms import PPO
from rsl_rl.modules import ActorCritic
from isaaclab.envs import DirectMARLEnv, 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
from isaaclab.utils.dict import class_to_dict
import berkeley_humanoid_lite.tasks
class RslRlVecEnvWrapperFixed(RslRlVecEnvWrapper):
def __init__(self, env):
self.env = env
base_env = env.unwrapped
self.unwrapped_env = base_env
self.device = base_env.device
obs_manager = base_env.observation_manager
def _to_int(val):
if isinstance(val, (tuple, list)): return int(val[0])
return int(val)
self.num_obs = _to_int(obs_manager.group_obs_dim["policy"])
self.num_actions = _to_int(base_env.action_manager.total_action_dim)
self.num_envs = base_env.num_envs
if "critic" in obs_manager.group_obs_dim:
self.num_privileged_obs = _to_int(obs_manager.group_obs_dim["critic"])
else:
self.num_privileged_obs = None
self.episode_length_buf = base_env.episode_length_buf
def _strip_tensordict(self, obs):
if obs is None: return None
if not isinstance(obs, torch.Tensor):
if isinstance(obs, dict):
return torch.cat(list(obs.values()), dim=-1)
return obs.view(obs.shape)
def get_observations(self):
obs = self.unwrapped.observation_manager.compute_group("policy")
return self._strip_tensordict(obs)
def get_privileged_observations(self):
if self.num_privileged_obs is not None:
obs = self.unwrapped.observation_manager.compute_group("critic")
return self._strip_tensordict(obs)
return None
def step(self, actions):
obs_dict, rew, terminated, truncated, extras = self.env.step(actions)
obs = self._strip_tensordict(obs_dict["policy"])
privileged_obs = None
if "critic" in obs_dict:
privileged_obs = self._strip_tensordict(obs_dict["critic"])
elif self.num_privileged_obs is not None:
privileged_obs = self.get_privileged_observations()
dones = terminated | truncated
return obs, privileged_obs, rew, dones, extras
def reset(self):
obs_dict, _ = self.env.reset()
obs = self._strip_tensordict(obs_dict["policy"])
return obs, self.get_privileged_observations()
def filter_dict(raw_dict, whitelist):
return {k: v for k, v in raw_dict.items() if k in whitelist}
PPO_WHITELIST = ['value_loss_coef', 'use_clipped_value_loss', 'clip_param', 'entropy_coef',
'num_learning_epochs', 'num_mini_batches', 'learning_rate', 'schedule',
'gamma', 'lam', 'desired_kl', 'max_grad_norm']
POLICY_WHITELIST = ['init_noise_std', 'actor_hidden_dims', 'critic_hidden_dims', 'activation']
@hydra_task_config(args_cli.task, "rsl_rl_cfg_entry_point")
def main(env_cfg, agent_cfg: RslRlOnPolicyRunnerCfg):
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 1
if hasattr(env_cfg.scene, "terrain"):
print("[INFO] Play模式: 强制地形为无限平面 (Plane)")
env_cfg.scene.terrain.terrain_type = "plane"
env_cfg.scene.terrain.terrain_generator = None
if hasattr(env_cfg.scene, "robot"):
env_cfg.scene.robot.init_state.pos = (0.0, 0.0, 0)
print("[INFO] Play模式: 强制初始高度为 1.05m")
log_root_path = os.path.join("logs", "rsl_rl", agent_cfg.experiment_name)
log_root_path = os.path.abspath(log_root_path)
load_run = agent_cfg.load_run
if load_run == "-1":
load_run = None
resume_path = None
try:
resume_path = get_checkpoint_path(log_root_path, load_run, agent_cfg.load_checkpoint)
print(f"[INFO] 加载模型路径: {resume_path}")
except Exception as e:
print(f"[ERROR] 无法找到模型 checkpoint: {e}")
simulation_app.close()
sys.exit(1)
print(f"[INFO] 正在创建环境: {args_cli.task}")
env = gym.make(args_cli.task, cfg=env_cfg)
if isinstance(env.unwrapped, DirectMARLEnv):
env = multi_agent_to_single_agent(env)
env = RslRlVecEnvWrapperFixed(env)
raw_dict = class_to_dict(agent_cfg) if not isinstance(agent_cfg, dict) else agent_cfg
rsl_cfg = {
"runner": {
"policy_class_name": "ActorCritic",
"algorithm_class_name": "PPO",
"experiment_name": agent_cfg.experiment_name,
"checkpoint": resume_path,
"num_steps_per_env": agent_cfg.num_steps_per_env,
"max_iterations": agent_cfg.max_iterations,
"save_interval": agent_cfg.save_interval,
"run_name": agent_cfg.run_name,
},
"algorithm": filter_dict(raw_dict.get("algorithm", {}), PPO_WHITELIST),
"policy": filter_dict(raw_dict.get("policy", {}), POLICY_WHITELIST),
}
runner = OnPolicyRunner(env, rsl_cfg, log_dir=None, device=agent_cfg.device)
runner.load(resume_path)
policy = runner.get_inference_policy(device=env.device)
print("-" * 80)
print("[INFO] 启动成功！在 Isaac Sim 中按 'F' 键跟随机器人。")
print("-" * 80)
obs, _ = env.reset()
while simulation_app.is_running():
with torch.inference_mode():
actions = policy(obs)
obs, _, _, _, _ = env.step(actions)
env.close()
if __name__ == "__main__":
main()
simulation_app.close()
|
