examples/record_trained_policy.py

"""Record a TRAINED Isaac Lab rsl_rl policy to a LeRobot v3 dataset through strands-robots.

Chain proven here (strands-robots feat/isaaclab-trainer @fa66fc68, PR #4227):
  train_policy(provider="isaaclab")  -> rsl_rl run dir with model_<N>.pt + params/{env,agent}.yaml
  THIS SCRIPT (runs in the Isaac Lab venv, strands on PYTHONPATH):
    1. rebuild the task env (play-mode cfg, same physics preset) + a strands camera per env
    2. load the rsl_rl checkpoint with OnPolicyRunner, export JIT + ONNX (Isaac Lab's own exporter)
    3. wrap the exported JIT actor as a strands `Policy` (RslRlJitPolicy) -- IL-X-006: strands'
       create_policy("rl") cannot load the rsl_rl ELU actor, so the example carries the adapter
    4. roll it out: every env i in 0..N-1 is one episode (from a common reset to its first done or
       --frames), actions come from policy.get_actions_sync() per env (strands single-robot API)
    5. record with strands `DatasetRecorder` (observation.state = joint pos + policy obs + root pose,
       action = policy action, observation.images.<cam> = RTX camera video, task string, fps)
  --verify ROOT (fresh process, no Kit): strands verify_dataset + LeRobotDataset load + video decode
  + frame counts + NaN scan.

Run (Isaac Lab venv):
  OMNI_KIT_ACCEPT_EULA=YES PYTHONPATH=$T/isaaclab-wt $ISAACLAB_PYTHON record_trained_policy.py \
     --task Isaac-Velocity-Rough-G1 --checkpoint RUN/model_1499.pt --episodes 12 --frames 500 \
     --override physics=newton_mjwarp --cam chase --root out/ds --repo_id cagataydev/x
  $ISAACLAB_PYTHON record_trained_policy.py --verify out/ds --repo_id cagataydev/x
"""

from __future__ import annotations

import argparse
import json
import os
import shutil
import sys
import time

ap = argparse.ArgumentParser()
ap.add_argument("--task")
ap.add_argument("--checkpoint")
ap.add_argument("--episodes", type=int, default=12, help="= num_envs; env i is episode i")
ap.add_argument("--frames", type=int, default=500)
ap.add_argument("--override", action="append", default=[], help="hydra override, e.g. physics=newton_mjwarp")
ap.add_argument("--cam", default="chase", choices=["chase", "fixed", "none"])
ap.add_argument("--cam_name", default=None)
ap.add_argument("--eye", default=None, help="eye offset x,y,z (chase: from root; fixed: from env origin)")
ap.add_argument("--target", default="0,0,0", help="fixed cam: look-at offset from env origin (or root with chase)")
ap.add_argument("--attach", default="Robot/pelvis", help="chase cam parent prim under the env ns (a robot link)")
ap.add_argument("--width", type=int, default=320)
ap.add_argument("--height", type=int, default=240)
ap.add_argument("--no_play_mode", action="store_true")
ap.add_argument("--task_str", default=None, help="language task string stored per frame")
ap.add_argument("--robot_type", default=None)
ap.add_argument("--root")
ap.add_argument("--repo_id", default="local/isaaclab_policy")
ap.add_argument("--min_len", type=int, default=1)
ap.add_argument("--verify", default=None, help="verify an existing dataset root and exit (no Kit)")
ap.add_argument("--seed", type=int, default=7)
a = ap.parse_args()
sys.argv = sys.argv[:1]


# ----------------------------------------------------------------------------- verify (no Kit)
def verify(root: str, repo_id: str, expected: int | None = None) -> dict:
    import numpy as np
    import pandas as pd

    out: dict = {"root": root}
    from strands_robots.verify_dataset import verify_dataset

    rep = verify_dataset(root, expected=expected, min_frames=1)
    out["strands_verify"] = rep
    from lerobot.datasets.lerobot_dataset import LeRobotDataset

    ds = LeRobotDataset(repo_id, root=root)
    out["num_episodes"], out["num_frames"], out["fps"] = ds.num_episodes, ds.num_frames, ds.fps
    info = json.load(open(os.path.join(root, "meta", "info.json")))
    out["features"] = {k: v.get("shape") for k, v in info["features"].items()}
    vid_keys = [k for k, v in info["features"].items() if v.get("dtype") == "video"]
    # decode first / middle / last frame through LeRobot (torchcodec/pyav)
    dec = {}
    for idx in (0, ds.num_frames // 2, ds.num_frames - 1):
        it = ds[idx]
        for k in vid_keys:
            t = it[k]
            dec[f"{k}@{idx}"] = {"shape": list(t.shape), "mean": round(float(t.float().mean()), 4)}
    out["video_decode"] = dec
    # NaN scan + per-episode frame counts straight from parquet
    files = sorted(
        os.path.join(dp, f) for dp, _, fs in os.walk(os.path.join(root, "data")) for f in fs if f.endswith(".parquet")
    )
    df = pd.concat([pd.read_parquet(f) for f in files])
    nan = 0
    for c in ("observation.state", "action"):
        arr = np.stack(df[c].to_numpy())
        nan += int((~np.isfinite(arr)).sum())
        out[f"{c}.dim"] = arr.shape[1]
    out["nan_or_inf"] = nan
    out["frames_per_episode"] = df.groupby("episode_index").size().tolist()
    out["tasks"] = sorted(set(ds.meta.tasks.index)) if hasattr(ds.meta.tasks, "index") else str(ds.meta.tasks)[:200]
    ok_verify = bool(rep.get("ok", rep.get("valid", True))) if isinstance(rep, dict) else True
    out["OK"] = bool(
        nan == 0
        and ds.num_frames == sum(out["frames_per_episode"])
        and (expected is None or ds.num_episodes == expected)
        and all(v["shape"][0] == 3 for v in dec.values())
        and ok_verify
    )
    return out


if a.verify:
    res = verify(a.verify, a.repo_id)
    print("[record] VERIFY " + json.dumps(res, default=str))
    json.dump(res, open(a.verify.rstrip("/") + "_verify.json", "w"), indent=1, default=str)
    sys.exit(0 if res["OK"] else 3)

# ----------------------------------------------------------------------------- record (Kit)
import gymnasium as gym  # noqa: E402
import numpy as np  # noqa: E402
import torch  # noqa: E402
from isaaclab.app import launch_simulation  # noqa: E402

import isaaclab_tasks  # noqa: E402,F401
from isaaclab_tasks.utils import resolve_task_config  # noqa: E402

N = a.episodes
env_cfg, agent_cfg = resolve_task_config(
    a.task, "rsl_rl_cfg_entry_point", play_mode=not a.no_play_mode, overrides=list(a.override)
)
env_cfg.scene.num_envs = N
env_cfg.sim.device = "cuda:0"
env_cfg.seed = a.seed
CAM = a.cam_name or ("chase" if a.cam == "chase" else "front")
use_cam = a.cam != "none"
if use_cam:
    import isaaclab.sim as sim_utils
    from isaaclab.sensors import CameraCfg

    import math

    def look_quat_world(eye, tgt):
        """xyzw quat, 'world' camera convention (+X forward, +Z up), looking from eye to tgt (no roll)."""
        dx, dy, dz = (t - e for t, e in zip(tgt, eye))
        yaw, pitch = math.atan2(dy, dx), math.atan2(-dz, math.hypot(dx, dy))
        sz, cz, sy, cy = math.sin(yaw / 2), math.cos(yaw / 2), math.sin(pitch / 2), math.cos(pitch / 2)
        return (-sz * sy, cz * sy, sz * cy, cz * cy)

    EYE = [float(v) for v in (a.eye or ("1.8,-2.6,0.6" if a.cam == "chase" else "1.2,0,0.8")).split(",")]
    TGT = [float(v) for v in a.target.split(",")]
    parent = "{ENV_REGEX_NS}/" + (a.attach if a.cam == "chase" else "")
    # NB: Camera.set_world_poses[_from_view] writes Fabric poses the RTX render does not pick up in
    # Isaac Lab 3.0.0rc1 (camera stays at its spawn pose) -> static offsets on a parent prim instead.
    setattr(
        env_cfg.scene,
        "strands_cam",
        CameraCfg(
            prim_path=parent.rstrip("/") + "/StrandsCam",
            update_period=0.0,
            height=a.height,
            width=a.width,
            data_types=["rgb"],
            spawn=sim_utils.PinholeCameraCfg(focal_length=18.0, clipping_range=(0.02, 50.0)),
            offset=CameraCfg.OffsetCfg(pos=tuple(EYE), rot=look_quat_world(EYE, TGT), convention="world"),
        ),
    )
largs = {"device": "cuda:0", "headless": True, "enable_cameras": use_cam}
meta: dict = {
    "task": a.task, "checkpoint": a.checkpoint, "episodes_requested": N, "frames": a.frames,
    "overrides": a.override, "cam": a.cam, "cam_name": CAM, "res": [a.width, a.height],
    "play_mode": not a.no_play_mode, "seed": a.seed, "cmd": " ".join(map(str, sys.orig_argv)),
}
t0 = time.time()

from strands_robots.policies.base import Policy  # noqa: E402


class RslRlJitPolicy(Policy):
    """strands Policy over an Isaac Lab-exported rsl_rl TorchScript actor (obs normalizer + ELU MLP).

    observation_dict must carry "policy_obs": the concatenated actor observation groups (1-D).
    Returns a 1-step chunk [{action_name: value}]. Adapter exists because create_policy("rl")
    rebuilds strands' Tanh actor and cannot load rsl_rl weights (IL-X-006).
    """

    def __init__(self, jit_path: str, action_names: list[str], device: str = "cpu"):
        self.net = torch.jit.load(jit_path, map_location=device).eval()
        self.device = device
        self.action_names = list(action_names)
        self.robot_state_keys: list[str] = []

    def set_robot_state_keys(self, robot_state_keys: list[str]) -> None:
        self.robot_state_keys = list(robot_state_keys)

    @property
    def provider_name(self) -> str:  # type: ignore[override]
        return "isaaclab_rsl_rl_jit"

    def reset(self) -> None:
        self.net.reset()

    async def get_actions(self, observation_dict, instruction, **kwargs):
        x = torch.as_tensor(np.asarray(observation_dict["policy_obs"], dtype=np.float32), device=self.device)
        with torch.inference_mode():
            y = self.net(x.unsqueeze(0))[0].cpu().numpy()
        return [{n: float(v) for n, v in zip(self.action_names, y)}]


with launch_simulation(env_cfg, largs):
    from rsl_rl.runners import OnPolicyRunner

    from isaaclab_rl.rsl_rl import RslRlVecEnvWrapper, handle_deprecated_rsl_rl_cfg
    import importlib.metadata as md

    agent_cfg = handle_deprecated_rsl_rl_cfg(agent_cfg, md.version("rsl-rl-lib"))
    genv = gym.make(a.task, cfg=env_cfg)
    env = RslRlVecEnvWrapper(genv, clip_actions=agent_cfg.clip_actions)
    u = env.unwrapped
    runner = OnPolicyRunner(env, agent_cfg.to_dict(), log_dir=None, device="cuda:0")
    runner.load(a.checkpoint)
    exp_dir = os.path.join(os.path.dirname(a.checkpoint), "exported")
    runner.export_policy_to_jit(path=exp_dir, filename="policy.pt")
    try:
        runner.export_policy_to_onnx(path=exp_dir, filename="policy.onnx")
    except Exception as e:  # noqa: BLE001
        meta["onnx_error"] = f"{type(e).__name__}: {e}"[:300]
    gpol = runner.alg.get_policy()
    obs_groups = list(getattr(gpol, "obs_groups", ["policy"]))
    meta["actor_obs_groups"] = obs_groups
    robot = next(iter(u.scene.articulations.values()))
    jn = list(robot.joint_names)
    adim = int(env.num_actions)
    names: list[str] = []
    am = getattr(u, "action_manager", None)
    if am is not None:
        for tname, term in am._terms.items():
            jj = list(getattr(term, "_joint_names", []) or [])
            names += [f"{tname}.{j}" for j in jj] if len(jj) == term.action_dim else [f"{tname}.{k}" for k in range(term.action_dim)]
    if len(names) != adim:
        idx = getattr(u, "actuated_dof_indices", None)
        names = [f"act.{jn[i]}" for i in idx] if idx is not None and len(idx) == adim else [f"a{k:02d}" for k in range(adim)]
    pol = RslRlJitPolicy(os.path.join(exp_dir, "policy.pt"), names)
    pol.set_robot_state_keys(jn)
    fps = int(round(1.0 / u.step_dt))
    meta.update(joint_names=jn, action_names=names, action_dim=adim, fps=fps, step_dt=u.step_dt,
                env_build_s=round(time.time() - t0, 1))

    obs = env.get_observations()
    pobs = torch.cat([obs[g] for g in obs_groups], dim=-1)
    D = int(pobs.shape[-1])
    meta["policy_obs_dim"] = D
    print(f"[record] env ready {meta['env_build_s']}s joints={len(jn)} act={adim} obs={D} fps={fps}")
    # consistency: strands-wrapped JIT vs rsl_rl inference policy on the same obs
    with torch.inference_mode():
        ref = runner.get_inference_policy(device="cuda:0")(obs)[0].cpu().numpy()
    got = np.array(list(pol.get_actions_sync({"policy_obs": pobs[0].cpu().numpy()}, "")[0].values()))
    meta["jit_vs_rslrl_max_abs"] = float(np.abs(ref - got).max())

    from strands_robots.dataset_recorder import DatasetRecorder

    task_str = a.task_str or f"{a.task}: trained rsl_rl PPO policy rollout"
    shutil.rmtree(a.root, ignore_errors=True)
    os.makedirs(os.path.dirname(os.path.abspath(a.root)), exist_ok=True)
    rec = DatasetRecorder.create(
        repo_id=a.repo_id, fps=fps, robot_type=a.robot_type or a.task, joint_names=jn, action_names=names,
        camera_keys=[CAM] if use_cam else None, camera_dims={CAM: (a.height, a.width)} if use_cam else None,
        extra_state_specs=[("policy_obs", [str(k) for k in range(D)]), ("root_pos", ["x", "y", "z"]),
                           ("root_quat", ["x", "y", "z", "w"])],
        task=task_str, root=a.root, use_videos=use_cam, overwrite=True,
    )
    def aim():
        pass

    buf: list[list] = [[] for _ in range(N)]  # per env: (obs_dict, act_dict)
    alive = np.ones(N, bool)
    ret = np.zeros(N)
    lens = np.zeros(N, int)
    term_reason = ["frames"] * N
    pol_s = step_s = 0.0
    pol.reset()
    aim()
    env.step(torch.zeros(N, adim, device=u.device))  # warm the renderer after aiming
    obs, _ = env.reset() if hasattr(env, "reset") else (env.get_observations(), None)
    obs = env.get_observations()
    for f in range(a.frames):
        aim()
        if f in (0, 1, 10) and use_cam:
            cd = u.scene["strands_cam"].data
            pw = cd.pos_w.torch if hasattr(cd.pos_w, "torch") else cd.pos_w
            rpw = robot.data.root_pos_w.torch if hasattr(robot.data.root_pos_w, "torch") else robot.data.root_pos_w
            meta.setdefault("debug", []).append({"f": f, "cam_pos_w0": pw[0].tolist(), "root_pos_w0": rpw[0].tolist()})
        pobs = torch.cat([obs[g] for g in obs_groups], dim=-1)
        pc = pobs.cpu().numpy()
        q = (robot.data.joint_pos.torch if hasattr(robot.data.joint_pos, "torch") else robot.data.joint_pos).cpu().numpy()
        rp = robot.data.root_pos_w
        rp = (rp.torch if hasattr(rp, "torch") else rp).cpu().numpy() - u.scene.env_origins.cpu().numpy()
        rq = robot.data.root_quat_w
        rq = (rq.torch if hasattr(rq, "torch") else rq).cpu().numpy()
        img = None
        if use_cam:
            im = u.scene["strands_cam"].data.output["rgb"]
            im = im.torch if hasattr(im, "torch") else im
            img = im[..., :3].to(torch.uint8).cpu().numpy()
        ts = time.time()
        acts = np.zeros((N, adim), np.float32)
        for i in range(N):
            ad = pol.get_actions_sync({"policy_obs": pc[i]}, task_str)[0]
            acts[i] = [ad[n] for n in names]
        pol_s += time.time() - ts
        for i in range(N):
            if not alive[i]:
                continue
            ob = {n: float(q[i, k]) for k, n in enumerate(jn)}
            ob.update(policy_obs=pc[i].astype(np.float32), root_pos=rp[i].astype(np.float32), root_quat=rq[i].astype(np.float32))
            if img is not None:
                ob[CAM] = img[i]
            buf[i].append((ob, {n: float(acts[i, k]) for k, n in enumerate(names)}))
            lens[i] += 1
        ts = time.time()
        with torch.inference_mode():
            obs, rew, dones, extras = env.step(torch.as_tensor(acts, device=u.device))
        step_s += time.time() - ts
        r = rew.cpu().numpy()
        d = dones.cpu().numpy().astype(bool)
        ret[alive] += r[alive]
        for i in np.where(d & alive)[0]:
            alive[i] = False
            to = extras.get("time_outs")
            term_reason[i] = "time_out" if (to is not None and bool(to[i])) else "terminated"
        if not alive.any():
            break
    print("[record] lens", lens.tolist(), term_reason, flush=True)
    meta.update(policy_ms_per_step_all_envs=round(1000 * pol_s / max(1, f + 1), 2),
                env_ms_per_step=round(1000 * step_s / max(1, f + 1), 2))
    # dump a 2x2 playback grid video (envs 0..3) + episode 0 video for model cards
    if use_cam:
        import imageio.v2 as iio

        base = a.root.rstrip("/")
        k = min(4, N)
        show = [int(i) for i in np.argsort(-lens, kind="stable")[:k]]  # the 4 longest episodes
        meta["playback_envs"] = show
        T = int(min(lens[show]))
        w = iio.get_writer(base + "_playback.mp4", fps=fps, codec="libx264", quality=7, macro_block_size=8)
        for t in range(T):
            ims = [buf[i][t][0][CAM] for i in show] + [np.zeros_like(buf[0][0][0][CAM])] * (4 - k)
            w.append_data(np.concatenate([np.concatenate(ims[:2], 1), np.concatenate(ims[2:4], 1)], 0))
        w.close()
        iio.imwrite(base + "_frame.png", buf[show[0]][min(T - 1, 60)][0][CAM])
        meta["cam_mean"] = float(np.mean([buf[i][-1][0][CAM].mean() for i in range(N) if buf[i]]))
    ts = time.time()
    kept = []
    for i in range(N):
        if lens[i] < a.min_len:
            continue
        for ob, ac in buf[i]:
            rec.add_frame(ob, ac, task=task_str, camera_keys=[CAM] if use_cam else None)
        rec.save_episode()
        kept.append(i)
        buf[i] = []
    rec.finalize()
    meta.update(record_s=round(time.time() - ts, 1), episodes_saved=len(kept), kept_envs=kept,
                episode_len=lens.tolist(), episode_return=[round(float(x), 3) for x in ret],
                episode_end=term_reason, mean_return=round(float(ret.mean()), 3),
                total_frames=int(lens[kept].sum()), total_wall_s=round(time.time() - t0, 1))
    print("[record] RESULT " + json.dumps(meta, default=str))
    json.dump(meta, open(a.root.rstrip("/") + "_record.json", "w"), indent=1, default=str)
    sys.stdout.flush()
    genv.close()  # Kit exits the process here; verify runs in a fresh process