#!/usr/bin/env python """ build_offline_buffer.py (Agent C) Convert the teleop demo dataset `banana_in_pot_lerobot` into the HIL-SERL OFFLINE DEMO BUFFER: a standard LeRobot v3 dataset that already carries the RL columns the SAC learner ingests via ReplayBuffer.from_lerobot_dataset: observation.state float32 [7] (ur_q1..6 + grip_pos) -- MATCHES online obs observation.images.cam1 video 3x128x128 (full-frame, resized) observation.images.cam2 video 3x128x128 action float32 [4] [delta_x, delta_y, delta_z, gripper] next.reward float32 [1] (Agent B success classifier, thr=0.7) next.done bool [1] (success onset OR episode end) WHY these choices (citations in hilserl/rl_dataset_schema.json + notes md): - action = base-frame TCP displacement / end_effector_step_sizes (Agent D rule). We derive the TCP position via FORWARD KINEMATICS of the dataset's OWN joint state (observation.state[:6]), NOT the raw-h5 tcp_pose. Reason: (a) the online deploy reference is `FK(current joints)` recomputed every step (EEReferenceAndDelta, use_latched_reference=False, gym_manipulator.py:507), so FK(joints) is exactly the quantity whose per-step delta the policy must reproduce; (b) it is perfectly frame-aligned to the images/state the policy observes (no async-stream matching, no episode->take mapping). Agent D validated FK == recorded tcp_pose to 0.85 mm at rest across all 51 takes, so this is equivalent to the recorded TCP. - gripper action = discrete class {0=close, 1=stay, 2=open} from grip_pos transitions (robot_kinematic_processor.py:408-412 semantics; grip_cmd is NaN-prone so grip_pos). - observation.state kept as the raw 7-d joint state: the online env base observation is `agent_pos` = motor-bus joint positions (gym_manipulator.py:173-177) with all ObservationConfig add_* flags defaulting False (envs/configs.py:266-268) -> 7-dim. - images 3x128x128 to match Agent B's reward classifier (trained at 128x128); the online env must set image_preprocessing.resize_size=[128,128] to match (documented). """ from __future__ import annotations import json import os import sys import numpy as np import torch import torch.nn.functional as F HERE = os.path.dirname(os.path.abspath(__file__)) ROOT = os.path.dirname(HERE) sys.path.insert(0, HERE) from lerobot.datasets.lerobot_dataset import LeRobotDataset # noqa: E402 from lerobot.rewards.classifier.modeling_classifier import Classifier # noqa: E402 # ---------------------------------------------------------------- config SRC_ROOT = os.path.join(ROOT, "banana_in_pot_lerobot") OUT_ROOT = os.path.join(HERE, "banana_rl_lerobot") REPO_ID = "theo/banana_in_pot_rl" CKPT = os.path.join(HERE, "reward_classifier", "checkpoint") STEP_SIZE = 0.05 # metres per unit; end_effector_step_sizes x=y=z SUCCESS_THRESHOLD = 0.7 # Agent B caveat SUCCESS_REWARD = 1.0 IMG = 128 GRIP_EPS = 0.03 # deadzone on Δgrip_pos for discrete gripper class GRIPPER_CLOSE, GRIPPER_STAY, GRIPPER_OPEN = 0.0, 1.0, 2.0 DEVICE = "cuda" if torch.cuda.is_available() else "cpu" # ---------------------------------------------------------------- UR7e FK # scipy-free re-implementation of Agent D's validated FK (hilserl/joint_to_ee.py): # same default_kinematics.yaml joint origins, same RPY(xyz, URDF/extrinsic) + Rz(theta) # chain shoulder->wrist_3, base_link frame, no tool offset (recorded TCP == flange). import yaml # noqa: E402 _KIN_YAML = "/opt/ros/jazzy/share/ur_description/config/ur7e/default_kinematics.yaml" _LINK_ORDER = ["shoulder", "upper_arm", "forearm", "wrist_1", "wrist_2", "wrist_3"] def _Rx(a): c, s = np.cos(a), np.sin(a); return np.array([[1, 0, 0], [0, c, -s], [0, s, c]]) def _Ry(a): c, s = np.cos(a), np.sin(a); return np.array([[c, 0, s], [0, 1, 0], [-s, 0, c]]) def _Rz(a): c, s = np.cos(a), np.sin(a); return np.array([[c, -s, 0], [s, c, 0], [0, 0, 1]]) def _rpyxyz_to_T(p): T = np.eye(4) T[:3, :3] = _Rz(p["yaw"]) @ _Ry(p["pitch"]) @ _Rx(p["roll"]) # URDF RPY = extrinsic xyz T[:3, 3] = [p["x"], p["y"], p["z"]] return T with open(_KIN_YAML) as _f: _KIN = yaml.safe_load(_f)["kinematics"] _FIXED = [_rpyxyz_to_T(_KIN[k]) for k in _LINK_ORDER] def _Rz4(a): T = np.eye(4); T[:3, :3] = _Rz(a); return T def fk_batch(q_rad_seq): """(N,6) joint angles [rad] -> (N,4,4) TCP pose in base_link frame.""" out = np.zeros((len(q_rad_seq), 4, 4)) for n, q in enumerate(np.asarray(q_rad_seq, dtype=float)): T = np.eye(4) for i in range(6): T = T @ _FIXED[i] @ _Rz4(q[i]) out[n] = T return out def resize_uint8_hwc(chw_float01: torch.Tensor) -> np.ndarray: """(3,H,W) float[0,1] -> (128,128,3) uint8 HWC (matches Agent B cache + LeRobot video convention).""" x = F.interpolate(chw_float01.unsqueeze(0), size=(IMG, IMG), mode="bilinear", align_corners=False) return (x[0].permute(1, 2, 0).clamp(0, 1).numpy() * 255).astype(np.uint8) def norm_cam(uint8_hwc: np.ndarray, mean: torch.Tensor, std: torch.Tensor) -> torch.Tensor: """(N,128,128,3) uint8 -> (N,3,128,128) MEAN_STD-normalized (exactly Agent B inference).""" t = torch.from_numpy(uint8_hwc.astype(np.float32) / 255.0).permute(0, 3, 1, 2) return (t - mean[None, :, None, None]) / std[None, :, None, None] def discrete_gripper(grip: np.ndarray) -> np.ndarray: """grip_pos (N,) -> discrete gripper command per frame (N,). grip_pos increases when closing (grasp ~0.55), decreases when opening (~0.01). Online: 0=close, 1=stay, 2=open. Command at frame t reproduces grip[t]->grip[t+1].""" n = len(grip) cls = np.full(n, GRIPPER_STAY, dtype=np.float32) dg = np.diff(grip) # (n-1,) cls[:-1] = np.where(dg > GRIP_EPS, GRIPPER_CLOSE, np.where(dg < -GRIP_EPS, GRIPPER_OPEN, GRIPPER_STAY)) return cls def main(): print(f"device={DEVICE} src={SRC_ROOT}\n") ds = LeRobotDataset(REPO_ID, root=SRC_ROOT) task_str = ds.meta.tasks.index[0] if hasattr(ds.meta.tasks, "index") else list(ds.meta.tasks)[0] print("task string:", task_str) # per-camera normalization stats (same source as Agent B) stats = json.load(open(os.path.join(SRC_ROOT, "meta", "stats.json"))) mean1 = torch.tensor(np.array(stats["observation.images.cam1"]["mean"]).reshape(3), dtype=torch.float32) std1 = torch.tensor(np.array(stats["observation.images.cam1"]["std"]).reshape(3), dtype=torch.float32) mean2 = torch.tensor(np.array(stats["observation.images.cam2"]["mean"]).reshape(3), dtype=torch.float32) std2 = torch.tensor(np.array(stats["observation.images.cam2"]["std"]).reshape(3), dtype=torch.float32) clf = Classifier.from_pretrained(CKPT).to(DEVICE).eval() # ---- output dataset schema (RL columns) ---- features = { "observation.state": {"dtype": "float32", "shape": [7], "names": ["ur_q1", "ur_q2", "ur_q3", "ur_q4", "ur_q5", "ur_q6", "grip_pos"]}, "observation.images.cam1": {"dtype": "video", "shape": [IMG, IMG, 3], "names": ["height", "width", "channels"]}, "observation.images.cam2": {"dtype": "video", "shape": [IMG, IMG, 3], "names": ["height", "width", "channels"]}, "action": {"dtype": "float32", "shape": [4], "names": ["delta_x", "delta_y", "delta_z", "gripper"]}, "next.reward": {"dtype": "float32", "shape": [1], "names": None}, "next.done": {"dtype": "bool", "shape": [1], "names": None}, } if os.path.exists(OUT_ROOT): import shutil shutil.rmtree(OUT_ROOT) out = LeRobotDataset.create(repo_id=REPO_ID, fps=ds.fps, root=OUT_ROOT, features=features, use_videos=True) # ---- group frames by episode (single ordered pass) ---- from_idx = ds.meta.episodes["dataset_from_index"] to_idx = ds.meta.episodes["dataset_to_index"] n_ep = ds.meta.total_episodes all_actions = [] grip_class_counts = np.zeros(3, dtype=int) n_success_ep = 0 reward_timeline_ep0 = None prob_timeline_ep0 = None for ep in range(n_ep): lo, hi = int(from_idx[ep]), int(to_idx[ep]) idxs = list(range(lo, hi)) n = len(idxs) states = np.zeros((n, 7), dtype=np.float32) c1 = np.zeros((n, IMG, IMG, 3), dtype=np.uint8) c2 = np.zeros((n, IMG, IMG, 3), dtype=np.uint8) for j, i in enumerate(idxs): f = ds[i] states[j] = f["observation.state"].numpy() c1[j] = resize_uint8_hwc(f["observation.images.cam1"]) c2[j] = resize_uint8_hwc(f["observation.images.cam2"]) # --- action: FK(joints) -> base-frame TCP; delta / step_size --- pos = fk_batch(states[:, :6])[:, :3, 3] # (n,3) metres actions = np.zeros((n, 4), dtype=np.float32) actions[:-1, :3] = np.diff(pos, axis=0) / STEP_SIZE actions[:, 3] = discrete_gripper(states[:, 6]) all_actions.append(actions) for cval in (GRIPPER_CLOSE, GRIPPER_STAY, GRIPPER_OPEN): grip_class_counts[int(cval)] += int((actions[:, 3] == cval).sum()) # --- reward: Agent B classifier over frames (batched) --- probs = np.zeros(n, dtype=np.float32) with torch.no_grad(): for s in range(0, n, 256): e = min(n, s + 256) b1 = norm_cam(c1[s:e], mean1, std1).to(DEVICE) b2 = norm_cam(c2[s:e], mean2, std2).to(DEVICE) probs[s:e] = clf.predict([b1, b2]).probabilities.cpu().numpy() reward = (probs > SUCCESS_THRESHOLD).astype(np.float32) # --- done: success onset OR episode end --- done = np.zeros(n, dtype=bool) succ = np.where(reward > 0.5)[0] if len(succ): n_success_ep += 1 done[succ[0]] = True # success onset done[-1] = True # episode end if ep == 0: reward_timeline_ep0 = reward.copy() prob_timeline_ep0 = probs.copy() for j in range(n): out.add_frame({ "observation.state": states[j], "observation.images.cam1": c1[j], "observation.images.cam2": c2[j], "action": actions[j], "next.reward": np.array([reward[j]], dtype=np.float32), "next.done": np.array([done[j]], dtype=bool), "task": task_str, }) out.save_episode() print(f"ep {ep:02d}: n={n:4d} success_onset={'-' if not len(succ) else succ[0]:>4} " f"reward_frames={int(reward.sum()):4d} |Δp|/step max={np.abs(actions[:,:3]).max():.3f}") out.finalize() # ---------------------------------------------------------- stats report A = np.concatenate(all_actions, axis=0) cont = A[:, :3] rep = { "total_frames": int(A.shape[0]), "n_episodes": int(n_ep), "n_episodes_with_success_onset": int(n_success_ep), "step_size_m": STEP_SIZE, "success_threshold": SUCCESS_THRESHOLD, "action_continuous_xyz": { "min": cont.min(0).tolist(), "max": cont.max(0).tolist(), "q01": np.quantile(cont, 0.01, axis=0).tolist(), "q50": np.quantile(cont, 0.50, axis=0).tolist(), "q99": np.quantile(cont, 0.99, axis=0).tolist(), "frac_abs_gt_1": float((np.abs(cont) > 1.0).mean()), }, "gripper_class_counts": {"close(0)": int(grip_class_counts[0]), "stay(1)": int(grip_class_counts[1]), "open(2)": int(grip_class_counts[2])}, } print("\n==== ACTION / REWARD STATS ====") print(json.dumps(rep, indent=2)) json.dump(rep, open(os.path.join(HERE, "action_reward_stats.json"), "w"), indent=2) # ---------------------------------------------------------- plots try: import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt fig, ax = plt.subplots(1, 3, figsize=(16, 4)) for d, name in zip(range(3), ["Δx", "Δy", "Δz"]): ax[0].hist(cont[:, d], bins=120, alpha=0.5, label=name) ax[0].axvline(-1, color="k", ls="--", lw=0.8); ax[0].axvline(1, color="k", ls="--", lw=0.8) ax[0].set_title(f"EE-delta action (÷ step={STEP_SIZE} m)"); ax[0].set_xlabel("normalized delta"); ax[0].legend() ax[1].bar(["close(0)", "stay(1)", "open(2)"], grip_class_counts, color=["#d62728", "#7f7f7f", "#2ca02c"]) ax[1].set_title("gripper discrete class counts") ax[2].plot(prob_timeline_ep0, label="classifier prob") ax[2].plot(reward_timeline_ep0, label="next.reward", lw=2) ax[2].axhline(SUCCESS_THRESHOLD, color="k", ls="--", lw=0.8, label=f"thr={SUCCESS_THRESHOLD}") ax[2].set_title("episode 0: reward timeline"); ax[2].set_xlabel("frame"); ax[2].legend() fig.tight_layout() fig.savefig(os.path.join(HERE, "rl_dataset_action_reward.png"), dpi=110) print("saved rl_dataset_action_reward.png") except Exception as e: print("plot skipped:", e) # ---------------------------------------------------------- schema json schema = { "repo_id": REPO_ID, "root": OUT_ROOT, "codebase_version": "v3.0", "fps": int(ds.fps), "total_episodes": int(n_ep), "total_frames": int(A.shape[0]), "rl_columns": { "action": {"dtype": "float32", "shape": [4], "names": ["delta_x", "delta_y", "delta_z", "gripper"], "semantics": "xyz = base-frame TCP delta / end_effector_step_sizes(=0.05 m); " "gripper = discrete class {0=close,1=stay,2=open}"}, "next.reward": {"dtype": "float32", "shape": [1], "semantics": f"Agent B success classifier prob>{SUCCESS_THRESHOLD} -> {SUCCESS_REWARD}"}, "next.done": {"dtype": "bool", "shape": [1], "semantics": "True at success onset AND at episode end"}, }, "state_keys_for_replay_buffer": ["observation.state", "observation.images.cam1", "observation.images.cam2"], "observation.state": {"dtype": "float32", "shape": [7], "names": ["ur_q1", "ur_q2", "ur_q3", "ur_q4", "ur_q5", "ur_q6", "grip_pos"], "note": "matches online gym_manipulator agent_pos (all ObservationConfig add_* False)"}, "observation.images.cam1": {"dtype": "video", "stored_shape_hwc": [IMG, IMG, 3], "decoded_shape_chw": [3, IMG, IMG], "range": "float[0,1]"}, "observation.images.cam2": {"dtype": "video", "stored_shape_hwc": [IMG, IMG, 3], "decoded_shape_chw": [3, IMG, IMG], "range": "float[0,1]"}, "policy_output_features": {"action": {"type": "ACTION", "shape": [4]}}, "policy_input_features": { "observation.state": {"type": "STATE", "shape": [7]}, "observation.images.cam1": {"type": "VISUAL", "shape": [3, IMG, IMG]}, "observation.images.cam2": {"type": "VISUAL", "shape": [3, IMG, IMG]}, }, "online_env_requirements": { "image_preprocessing.resize_size": [IMG, IMG], "inverse_kinematics.end_effector_step_sizes": {"x": STEP_SIZE, "y": STEP_SIZE, "z": STEP_SIZE}, "policy.num_discrete_actions": 3, "reward_classifier.success_threshold": SUCCESS_THRESHOLD, "note": "crop_dataset_roi may be re-applied online; if ROI crop is used, re-run it on THIS " "dataset too so offline images match the online cropped size.", }, } json.dump(schema, open(os.path.join(HERE, "rl_dataset_schema.json"), "w"), indent=2) print("saved rl_dataset_schema.json") if __name__ == "__main__": main()