"""Study D: evaluate one frozen rsl_rl Franka lift policy in a chosen backend. Play-mode configuration (full curriculum difficulty, so full gravity; no observation noise; datasheet gripper speed), deterministic actions (actor mean, checkpoint normalization) and a fixed seed. One 12 s episode per environment. Conditions: play Isaac Lab's play-mode configuration unchanged: randomization on; starts from the task's validity-checked state bank, a quarter of them with the object already in the hand, drawn to keep success near 50% (the reference for the training metric). standard as play, but no spawn-in-hand starts and uniform draws from the state bank: every episode starts with the object free; the bank's collision checks still reject invalid starts. nominal as standard, with startup randomization off (friction, mass, inertia, joint gains and friction): both backends simulate the same nominal robot and objects. (An earlier version replaced the state bank with its raw reset terms; about one start in six then interpenetrated and ended on the first step. Those results are not used.) Matched starts: the backends draw different reset samples from the same seed. --save_starts FILE resets once (standard or nominal condition) and saves every env's start: arm and gripper joints, object pose and the first target. --load_starts FILE replaces the state bank with a reset term that writes exactly those states, inside the environment's normal reset (as Isaac Lab's own reset terms do); the first target is then set and the observations recomputed. Use with --single_target, which holds that target for the episode, so both engines face identical episodes. Success is the task's own sticky criterion: the object within 5 cm of a commanded target while thumb and finger touch it, at least once during the episode. Joint and body order must be identical in both backends for this task (Isaac Lab's transfer guide); the script verifies the joint names instead of assuming it. """ import argparse import json import sys import warp as wp wp.config.enable_backward = False import gymnasium as gym # noqa: E402 import numpy as np # noqa: E402 import torch # noqa: E402 from isaaclab.app import add_launcher_args, launch_simulation # noqa: E402 import isaaclab_tasks # noqa: E402,F401 from isaaclab_tasks.utils import resolve_task_config, setup_preset_cli # noqa: E402 STARTUP_RANDOMIZATION = ('robot_physics_material', 'object_physics_material', 'object_physics_inertia', 'joint_stiffness_and_damping', 'joint_friction', 'object_scale_mass', 'gripper_closing_speed') FINGER_SENSORS = ('panda_leftfinger_object_s', 'panda_rightfinger_object_s') def t(x): """Torch view of Isaac Lab data (ProxyArray, warp array or tensor).""" if hasattr(x, 'torch') and not isinstance(x, torch.Tensor): return x.torch return x if isinstance(x, torch.Tensor) else wp.to_torch(x) def configure(env_cfg, condition, num_envs, seed, load_starts=None): # play_mode() was already applied by resolve_task_config(play_mode=True). env_cfg.scene.num_envs = num_envs env_cfg.seed = seed ev = env_cfg.events if env_cfg.curriculum is not None: adr = env_cfg.curriculum.adr.params if adr['init_difficulty'] != adr['max_difficulty']: raise ValueError('play_mode did not set the curriculum to its maximum difficulty') if condition == 'play': return if condition == 'nominal': for name in STARTUP_RANDOMIZATION: setattr(ev, name, None) params = ev.conditional_reset.params params['terms']['reset_object_to_target'].params['probability'] = 0.0 params['success_monitor'] = None if load_starts: from isaaclab.managers import EventTermCfg ev.conditional_reset = None ev.reset_from_file = EventTermCfg(func=reset_from_file, mode='reset', params=dict(path=load_starts)) _STARTS = {} def reset_from_file(env, env_ids, path): """Reset term: write the saved start state of each env (joints, object pose; zero velocities).""" if path not in _STARTS: s = np.load(path) _STARTS[path] = {k: torch.tensor(s[k], device=env.device, dtype=torch.float32) for k in ('joint_pos', 'object_pose_env', 'command_b')} _STARTS[path]['joint_names'] = list(s['joint_names']) s = _STARTS[path] robot, obj = env.scene['robot'], env.scene['object'] if s['joint_names'] != list(robot.joint_names): raise ValueError('joint order differs from the saved starts') env_ids = torch.as_tensor(env_ids, device=env.device) q = s['joint_pos'][env_ids] robot.write_joint_position_to_sim_index(position=q, env_ids=env_ids) robot.write_joint_velocity_to_sim_index(velocity=torch.zeros_like(q), env_ids=env_ids) pose = s['object_pose_env'][env_ids].clone() pose[:, :3] += t(env.scene.env_origins)[env_ids] obj.write_root_pose_to_sim_index(root_pose=pose, env_ids=env_ids) obj.write_root_velocity_to_sim_index(root_velocity=torch.zeros(len(env_ids), 6, device=env.device), env_ids=env_ids) def save_starts(u, path): robot, obj = u.scene['robot'], u.scene['object'] pose = torch.cat([t(obj.data.root_pos_w)-t(u.scene.env_origins), t(obj.data.root_quat_w)], dim=1) np.savez(path, joint_pos=t(robot.data.joint_pos).cpu().numpy(), object_pose_env=pose.cpu().numpy(), command_b=u.command_manager.get_term('object_pose').pose_command_b.cpu().numpy(), joint_names=np.array(robot.joint_names)) def set_first_target(u, path): """After the reset: use the saved first target and rebuild the observation history with it.""" cmd = u.command_manager.get_term('object_pose') cmd.pose_command_b[:] = torch.tensor(np.load(path)['command_b'], device=u.device, dtype=torch.float32) u.observation_manager.reset() u.obs_buf = u.observation_manager.compute() def object_hulls(u, envs, num_points=600): """Convex hull of each recorded env's object, from points sampled on its actual mesh (object frame).""" from scipy.spatial import ConvexHull from isaaclab_tasks.core.lift.mdp.utils import sample_object_point_cloud pts = sample_object_point_cloud(u.num_envs, num_points, u.scene['object'].cfg.prim_path, device=u.device) hulls = {} for e in envs: p = pts[e].detach().cpu().double().numpy() h = ConvexHull(p) used = sorted(set(h.simplices.flatten().tolist())) index = {v: i for i, v in enumerate(used)} hulls[str(e)] = dict(vertices=p[used].round(6).tolist(), faces=[[index[int(v)] for v in f] for f in h.simplices], extent_m=(p.max(0)-p.min(0)).round(4).tolist()) return hulls def run(env, policy, wrapped, record_envs=()): from isaaclab.utils.math import combine_frame_transforms u = env.unwrapped robot, obj = u.scene['robot'], u.scene['object'] success_term = u.reward_manager.get_term_cfg('success').func sensors = [u.scene.sensors[s] for s in FINGER_SENSORS] arm = [i for i, n in enumerate(robot.joint_names) if n.startswith('panda_joint')] steps = int(round(u.cfg.episode_length_s/u.step_dt)) n = u.num_envs alive = torch.ones(n, dtype=torch.bool, device=u.device) success = torch.zeros(n, dtype=torch.bool, device=u.device) first_success = torch.full((n,), float('nan'), device=u.device) ended = torch.full((n,), steps*u.step_dt, device=u.device) cause = ['time_out']*n held = torch.zeros(n, device=u.device) count = torch.zeros(n, device=u.device) torque_sq = torch.zeros(n, device=u.device) peak_grip = torch.zeros(n, device=u.device) max_height = torch.full((n,), -1e9, device=u.device) start_height = t(obj.data.root_pos_w)[:, 2].clone() recording = {e: [] for e in record_envs} action_term = u.action_manager.get_term(u.action_manager.active_terms[0]) def record(time_s): # Isaac Lab 3 quaternions are (x, y, z, w); the replay viewer uses (w, x, y, z). Positions are env-local. origin = t(u.scene.env_origins) pose = t(robot.data.body_link_pose_w) opos, oquat = t(obj.data.root_pos_w), t(obj.data.root_quat_w) cmd = u.command_manager.get_command('object_pose') target, _ = combine_frame_transforms(t(robot.data.root_pos_w), t(robot.data.root_quat_w), cmd[:, :3]) finger = [t(s.data.net_forces_w).reshape(n, -1, 3).sum(dim=1) for s in sensors] for e in record_envs: if not alive[e]: continue o = origin[e] q = pose[e].detach().cpu().double() q[:, :3] -= o.cpu().double() q = q.numpy() recording[e].append(dict( time_s=round(time_s, 6), robot=np.c_[q[:, :3], q[:, 6:7], q[:, 3:6]].round(6).tolist(), object=[*(opos[e]-o).tolist(), float(oquat[e, 3]), *oquat[e, :3].tolist()], target=(target[e]-o).tolist(), finger_force_N=[f[e].tolist() for f in finger], joint_pos=t(robot.data.joint_pos)[e].tolist(), joint_target=t(action_term.processed_actions)[e].tolist(), succeeded=bool(success_term.succeeded[e]))) record(0.) obs = wrapped.get_observations() for k in range(steps): with torch.inference_mode(): obs, _, dones, _ = wrapped.step(policy(obs)) policy.reset(dones) m = alive.clone() # Sticky task flag; an env that terminates this step was already reset inside step(), so a success # reached on its very last step is not counted. now = success_term.succeeded.clone() newly = now & m & ~success first_success[newly] = (k+1)*u.step_dt success |= now & m forces = torch.stack([torch.linalg.norm(t(s.data.net_forces_w).reshape(n, -1, 3), dim=-1).amax(dim=1) for s in sensors], 1) gripped = (forces > 0.01).all(dim=1) held += m.float()*gripped.float() peak_grip = torch.where(m, torch.maximum(peak_grip, forces.amax(dim=1)), peak_grip) torque_sq += m.float()*(t(robot.data.applied_torque)[:, arm]**2).mean(dim=1) max_height = torch.where(m, torch.maximum(max_height, t(obj.data.root_pos_w)[:, 2]), max_height) count += m.float() died = dones.bool() & alive if died.any(): tm = u.termination_manager for name in tm.active_terms: hit = (tm.get_term(name).bool() & died).nonzero().flatten().tolist() for e in hit: cause[e] = name ended[died] = (k+1)*u.step_dt alive &= ~died record((k+1)*u.step_dt) c = count.clamp(min=1) per_env = dict(success=success, first_success_s=first_success, episode_s=ended, held_fraction=held/c, arm_torque_rms_Nm=torch.sqrt(torque_sq/c), peak_finger_force_N=peak_grip, lift_height_m=max_height-start_height) per_env = {k: v.detach().cpu().numpy().astype(float) for k, v in per_env.items()} causes = {name: cause.count(name) for name in sorted(set(cause))} summary = dict(success_rate=float(per_env['success'].mean()), n=n, terminations=causes, **{k: dict(mean=float(np.nanmean(v)), median=float(np.nanmedian(v))) for k, v in per_env.items() if k != 'success'}) rec = None if record_envs: rec = dict(body_names=list(robot.body_names), joint_names=list(robot.joint_names), dt_s=u.step_dt, object_hulls=object_hulls(u, record_envs), table=dict(pose=(t(u.scene['table'].data.root_pos_w)[record_envs[0]]-t(u.scene.env_origins)[record_envs[0]]).tolist(), size_m=list(u.cfg.scene.table.spawn.size)), envs={str(e): dict(frames=recording[e], success=bool(success[e]), termination=cause[e], first_success_s=None if torch.isnan(first_success[e]) else float(first_success[e])) for e in record_envs}) return summary, per_env, cause, rec def main(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument('--task', default='Isaac-Lift-Franka') parser.add_argument('--checkpoint', required=True) parser.add_argument('--condition', choices=('play', 'standard', 'nominal'), default='standard') parser.add_argument('--num_envs', type=int, default=512) parser.add_argument('--seed', type=int, default=1000) parser.add_argument('--output', required=True) parser.add_argument('--save_starts', help='Reset once, save every env start to this file and exit') parser.add_argument('--load_starts', help='Start every env from this saved file') parser.add_argument('--single_target', action='store_true', help='Hold the first target for the whole episode') parser.add_argument('--record_envs', type=int, nargs='*', default=[], help='Record full state of these envs (30 Hz)') add_launcher_args(parser) parser.set_defaults(device=None) args, hydra_args = setup_preset_cli(parser, None) sys.argv = [sys.argv[0]]+hydra_args from rsl_rl.runners import OnPolicyRunner import importlib.metadata from isaaclab_rl.rsl_rl import RslRlVecEnvWrapper, handle_deprecated_rsl_rl_cfg # The agent cfg is a PresetCfg here; resolving through Hydra selects the default runner cfg. env_cfg, agent_cfg = resolve_task_config(args.task, 'rsl_rl_cfg_entry_point', play_mode=True) configure(env_cfg, args.condition, args.num_envs, args.seed, args.load_starts) if args.single_target: env_cfg.commands.object_pose.resampling_time_range = (1e6, 1e6) args.device = env_cfg.sim.device agent_cfg = handle_deprecated_rsl_rl_cfg(agent_cfg, importlib.metadata.version('rsl-rl-lib')) out = dict(task=args.task, checkpoint=args.checkpoint, condition=args.condition, hydra=hydra_args, seed=args.seed, num_envs=args.num_envs, physics_class=type(env_cfg.sim.physics).__name__) with launch_simulation(env_cfg, args): env = gym.make(args.task, cfg=env_cfg) wrapped = RslRlVecEnvWrapper(env, clip_actions=agent_cfg.clip_actions) runner = OnPolicyRunner(wrapped, agent_cfg.to_dict(), log_dir=None, device=agent_cfg.device) runner.load(args.checkpoint) policy = runner.get_inference_policy(device=env.unwrapped.device) u = env.unwrapped out.update(joint_names=list(u.scene['robot'].joint_names), body_names=list(u.scene['robot'].body_names), events=sorted(k for k, v in vars(env_cfg.events).items() if v is not None)) torch.manual_seed(args.seed) wrapped.reset() u.episode_length_buf[:] = 0 if args.save_starts: save_starts(u, args.save_starts) print('STARTS_SAVED '+args.save_starts, flush=True) env.close() return if args.load_starts: set_first_target(u, args.load_starts) out.update(load_starts=args.load_starts, single_target=args.single_target) summary, per_env, causes, rec = run(env, policy, wrapped, args.record_envs) if rec is not None: import gzip with gzip.open(args.output.replace('.json', '-recording.json.gz'), 'xt') as f: json.dump(rec, f, separators=(',', ':')) out['result'] = summary # Curriculum-written gravity actually used (play mode fixes the curriculum at maximum difficulty). out['gravity'] = [list(map(float, g)) for g in u.event_manager.get_term_cfg('variable_gravity').params['gravity_distribution_params']] cm = u.curriculum_manager out['adr_difficulty_frac'] = float(cm._term_cfgs[cm.active_terms.index('adr')].func.difficulty_frac) out['termination_cause'] = causes print(json.dumps({k: v for k, v in summary.items()}), flush=True) with open(args.output, 'x') as f: json.dump(out, f, indent=1) np.savez_compressed(args.output.replace('.json', '.npz'), **per_env) print('EVAL_WRITTEN '+args.output, flush=True) env.close() if __name__ == '__main__': main()