build_offline_buffer.py

#!/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()