geonmin-kim/SO101-lv4-3color-cube-mat-to-mat-no-human-reset-alldir-v1-leadin-lookup

# option A: pandas
import pandas as pd
from huggingface_hub import hf_hub_download

REPO = "geonmin-kim/SO101-lv4-3color-cube-mat-to-mat-no-human-reset-alldir-v1-leadin-lookup"
path = hf_hub_download(REPO, "data/leadin.parquet", repo_type="dataset")
df = pd.read_parquet(path)

# option B: datasets
from datasets import load_dataset
ds = load_dataset(REPO, split="train")

def start_positions(df, color: str):
    """All start positions (and episode ids) of the given cube color."""
    sub = df[df.target_cube_color == color]
    return sub[["episode_index", "target_cube_start_x", "target_cube_start_y"]].to_numpy()

xy_red = start_positions(df, "red")   # shape (N, 3): [episode_index, x, y]

import numpy as np

def nearest_episode(df, color: str, x: float, y: float):
    """Episode whose start position is closest (Euclidean, px) to (x, y),
    among episodes with the same target color."""
    sub = df[df.target_cube_color == color].reset_index(drop=True)
    d = np.hypot(sub.target_cube_start_x - x, sub.target_cube_start_y - y)
    row = sub.iloc[int(d.idxmin())]
    return row, float(d.min())

row, dist = nearest_episode(df, "red", 480.0, 300.0)
print(int(row.episode_index), dist)

from scipy.spatial import cKDTree

trees = {
    c: (sub.reset_index(drop=True),
        cKDTree(sub[["target_cube_start_x", "target_cube_start_y"]].to_numpy()))
    for c, sub in df.groupby("target_cube_color")
}
sub, tree = trees["red"]
dist, idx = tree.query([480.0, 300.0])
row = sub.iloc[int(idx)]

def leadin_action_for(df, color: str, x: float, y: float):
    """Motor positions (6) of the lead-in pose whose episode's cube start
    position is nearest to (x, y)."""
    row, dist = nearest_episode(df, color, x, y)
    return np.asarray(row.reach_done_action, dtype=np.float32), int(row.episode_index), dist

action, ep_idx, dist = leadin_action_for(df, "red", 480.0, 300.0)
# action = [shoulder_pan, shoulder_lift, elbow_flex, wrist_flex, wrist_roll, gripper]