offline_eval.py

"""
Kiểm tra offline: so sánh action policy dự đoán với action THẬT đã ghi trong
chính dataset train. Nếu ngay trên dữ liệu train mà sai số vẫn lớn -> vấn đề
nằm ở quá trình train (chưa hội tụ / thiếu data), KHÔNG phải ở script test.

Cách chạy:
    python offline_eval.py --checkpoint <path_to_pretrained_model> --episode 0
"""
import argparse
import numpy as np
import torch

from lerobot.datasets.lerobot_dataset import LeRobotDataset
from lerobot.policies.act.modeling_act import ACTPolicy


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint", type=str, required=True)
    parser.add_argument("--dataset_root", type=str, default="./dataset_lerobot")
    parser.add_argument("--repo_id", type=str, default="local/vla_car_robot")
    parser.add_argument("--episode", type=int, default=0, help="Episode index để test (0..N-1)")
    parser.add_argument("--device", type=str, default="cuda")
    args = parser.parse_args()

    device = args.device if torch.cuda.is_available() else "cpu"

    print(f"[INFO] Loading dataset '{args.repo_id}' từ {args.dataset_root} ...")
    ds = LeRobotDataset(args.repo_id, root=args.dataset_root)

    print(f"[INFO] Loading policy từ {args.checkpoint} ...")
    policy = ACTPolicy.from_pretrained(args.checkpoint)
    policy.eval()
    policy.to(device)
    policy.reset()

    # Lấy các frame index thuộc đúng episode được chọn
    ep_mask = ds.hf_dataset["episode_index"]
    frame_indices = [i for i, e in enumerate(ep_mask) if int(e) == args.episode]
    if not frame_indices:
        print(f"[ERR] Không tìm thấy episode {args.episode} trong dataset.")
        return
    print(f"[INFO] Episode {args.episode} có {len(frame_indices)} frames. Đang chạy so sánh...")

    errors = []
    for idx in frame_indices:
        sample = ds[idx]
        obs = {
            "observation.state": sample["observation.state"].unsqueeze(0).to(device),
            "observation.images.overhead": sample["observation.images.overhead"].unsqueeze(0).to(device),
            "observation.images.gripper": sample["observation.images.gripper"].unsqueeze(0).to(device),
            "task": [sample.get("task", "")],
        }
        true_action = sample["action"].numpy()

        with torch.inference_mode():
            pred_action = policy.select_action(obs).squeeze(0).cpu().numpy()

        err = np.abs(pred_action - true_action)
        errors.append(err)
        print(f"  frame {idx:5d} | true={true_action} | pred={pred_action} | abs_err={err}")

    errors = np.array(errors)
    print("\n===== TỔNG KẾT =====")
    print(f"Mean abs error mỗi chiều action: {errors.mean(axis=0)}")
    print(f"Max abs error mỗi chiều action : {errors.max(axis=0)}")
    print(
        "\nNếu mean abs error CAO (ví dụ > 0.1-0.2 trong khi action nằm trong [-1,1]) "
        "ngay trên chính dữ liệu train -> model CHƯA HỌC TỐT (do thiếu epoch/data), "
        "không phải lỗi ở script test/inference."
    )


if __name__ == "__main__":
    main()