IDM v2 — Action EEF + 16-dim State 동시 예측 인프라 구현

TL;DR

설계 원칙: v1 코드 한 줄도 수정하지 않고 v2 라인업 분리. action_dim=32 그대로 — [0:12] action EEF, [12:28] state 16-dim, [28:32] unused (mask=False)
HF v3.0 schema 1:1 매핑: base_to_eef_pos(3) + base_to_eef_quat(4) + gripper_qpos(2) + base_pos(3) + base_quat(4). 추후 state_norm[:, 12:28] slice로 바로 observation.state 생성 가능
8개 신규 파일: precompute 스크립트, dataset v2, training v2, dump v2, smoke/본학습 shell 포함. 37,420 샘플, quat ‖q‖=1.000000 검증 완료
xyzw convention 확정: HDF5 검증으로 확정 (wxyz이면 base 방향 비현실적). 연속 frame sign canonicalize 적용
cv2 headless 버그 수정: opencv-python + opencv-python-headless 충돌 → headless 단일 설치로 해결. v1 학습의 latent 버그도 동시 수정

신규 파일

37,420

Dataset 샘플

1.000000

Quat ‖q‖

315M

Params (v1 동일)

Running

Smoke job 101

1 배경 / 목적

현재 IDM v1 출력 = action_eef[12] (Δxyz + Δrot + gripper_close + Δbase + control_mode)만. DreamGen 파이프라인과 GR00T N1 fine-tune에서 절대 pose/state (eef 절대 위치, base 위치, gripper qpos)가 필요한 상황.

기존 한계: action_eef Δ를 적분하면 drift 발생. 시뮬레이터 replay 없이는 절대 state 불가. HF DAVIAN-Robotics/robocasa-H50 v3.0은 observation.state[16]을 필수로 요구.

따라서 IDM의 출력을 확장하여 16-dim state도 동시 예측 — 단 v1과 v2 라인업을 완전 분리하여 기존 학습·평가 코드가 깨지지 않도록 설계.

2 작업 내용

설계 핵심: action_dim=32 그대로 유지

Target layout (32-dim): [0:12] action_eef (Δxyz, Δrot, gripper_close, Δbase, ctrl_mode) ← v1 그대로 [12:15] base_to_eef_pos (3) ← HF base-relative (world 아님 — 모바일 base 불변) [15:19] base_to_eef_quat (4) ← HF base-relative [19:21] gripper_qpos (2) ← HF (Panda finger ±0.04m) [21:24] base_pos (3) ← HF world frame [24:28] base_quat (4) ← HF world frame [28:32] unused (zero, mask=False) HF STATE_NAMES 순서: scripts/convert_to_lerobot_v3.py:48-56 1:1 매핑 → 추후 state_norm[:, 12:28] slice = observation.state 직결

Quaternion convention: xyzw 확정

HDF5 검증: base_quat[0] = [0, 0, 0.7071, 0.7071] - xyzw (w=last): rot([1,0,0]) → [0, 1, 0] → 90° around z ✓ (베이스 +y 방향) - wxyz (w=first): rot([1,0,0]) → [-1, 0, 0] → 비현실적 ✗ → xyzw 확정. HF v3.0도 HDF5 그대로 packing → xyzw 유지. Quat sign canonicalize (precompute 단): dot(q[t], q[t-1]) < 0 → q[t] = -q[t] (연속 frame discontinuity 방지) 추론 후 처리: quat → q /= ‖q‖ (unit norm 강제)

정규화 정책

영역	정책	이유
`[0:12]` action_eef	기존 stats.json min/max → [-1,1]	v1 그대로
`[12:15]` base_to_eef_pos	min/max → [-1,1]	linear
`[15:19]` base_to_eef_quat	raw (no scaling)	unit norm 보존, 각 dim 이미 [-1,1]
`[19:21]` gripper_qpos	min/max → [-1,1]	linear
`[21:24]` base_pos	min/max → [-1,1]	linear
`[24:28]` base_quat	raw (no scaling)	unit norm 보존

새 파일 목록 (v1 무수정)

파일	역할	검증
`scripts/build_train_mapping.py`	MG-30 train ep_idx → HDF5 demo 매핑 JSON	720/720 매핑, length sanity ✓
`scripts/precompute_idm_state.py`	HDF5 obs 5종 → idm_state/episode_*.npz (T,16) + state_stats.json	1440 ep, quat ‖q‖=1.000000 ✓
`gr00t/data/dataset_latent_idm_v2.py`	`RobocasaLatentIDMDatasetCachedV2` + collate_v2. 샘플 enumeration v1 동일, target [12:28] state packing 추가	37,420 샘플, shape/dtype/mask/quat-norm ✓
`IDM_dump/base_latent_A_cached_v2.yaml`	v1 yaml copy (architecture 동일)	OmegaConf load ✓
`scripts/idm_training_v2.py`	항상 V2 dataset/yaml 강제. v1 training.py 무수정 유지	CLI parse + 4 GPU DDP load ✓
`scripts/dump_idm_actions_latent_v2.py`	추론 → dim별 denorm ([0:12] action, [12:28] state), quat L2-norm, npz 저장	AST + 함수 정의 ✓
`scripts/run_idm_train_robocasa_latent_A_v2_smoke.sh`	4 GPU 30 step smoke	job 101 running (worker-5)
`scripts/run_idm_train_robocasa_latent_A_v2_60k_8gpu.sh`	8 GPU 60K 본 학습	smoke 통과 후 제출 예정

cv2 headless 버그 수정 (환경)

# 문제: conda env gr00t-idm에 opencv-python + opencv-python-headless 둘 다 설치됨 # → import 시 GUI 버전이 우선 → libgthread-2.0.so.0 의존 → worker에 없음 # → smoke job 95 (worker-5)에서 실패 # 수정: pip uninstall opencv-python pip install opencv-python-headless==4.13.0.92 # 확인: gr00t/utils/video.py는 VideoCapture, resize만 사용 → headless 호환 OK # 부수 효과: v1 학습의 잠재 버그도 동시 수정

3 결과

Precompute 결과

206,570

Train frames

720 ep

Train precompute

720 ep

Heldout precompute

1.000000

Quat ‖q‖ (train)

State Stats Sanity (train 720 ep)

field	dim	min	max	해석
base_to_eef_pos	3	[-0.41, -0.78, +0.02]	[+0.85, +0.80, +1.10]	Panda reach (m), 정상
base_to_eef_quat	4	[-0.38, -0.96, -0.55, -0.95]	[+1.00, +0.95, +0.96, +0.87]	xyzw raw [-1,1]
gripper_qpos	2	[-0.023, -0.046]	[+0.044, +0.013]	Panda finger ±0.04m, 대칭 mirror
base_pos	3	[+0.34, -5.08, +0.70]	[+5.60, -0.50, +0.73]	다양한 kitchen layout
base_quat	4	[-0.0, -0.0, -0.028, -1.00]	[-0.0, -0.0, +1.00, +1.00]	x,y≈0 → 모바일 base z축 회전만

Dataset v2 검증

37,420

V2 샘플 수

-80

v1 대비 (37,500)

315M

Params

Running

Smoke (job 101)

샘플 -80 원인: pooled cache 없는 4개 window skip (v1과 동일 이유 — 3개 length<93, 1개 손상 mp4). Dataset 동작 정상.

데이터 파일 위치

항목	경로
Train mapping	`outputs/mapping_mg30_train_to_hdf5.json`
Train state npz	`/home/nas_main/$USER/Dataset/robocasa/idm_mg30/idm_state/episode_*.npz`
Train state stats	`/home/nas_main/$USER/Dataset/robocasa/idm_mg30/meta/state_stats.json`
Heldout state npz	`/home/nas_main/$USER/Dataset/robocasa/idm_heldout/idm_state/episode_*.npz`
Smoke output	`outputs/idm_robocasa_latent_A_v2_smoke/`
60K output (예정)	`outputs/idm_robocasa_latent_A_v2_60k_8gpu/`

4 Takeaway

DreamGen 파이프라인 완성을 위한 핵심 인프라

v2 dump의 predicted_* 키들이 HF v3.0 observation.state 16-dim 1:1 매핑됨. 추후 convert_to_lerobot_v3.py에서 state_norm[:, 12:28] slice만으로 obs.state 생성 가능 — 추가 변환 코드 불필요.

모델 capacity 추가 부담 없음: architecture 변경 없이 output dim 32 그대로 유지. action [0:12] + state [12:28] 두 영역 동시 supervision이지만 실질적 추가 비용은 precompute 1440 ep npz 생성 (~5-10분) 뿐.

v1 ckpt 재사용 불가: [12:28] 영역이 v1에서 zero-supervision이었으므로 의미 있는 init이 아님. 60K from-scratch 재학습 필요 (~12-13h, 8 GPU). action mm error는 v1 수준 유지 예상 (capacity 공유이지만 영역 충분).

5 Next Steps

Smoke job 101 결과 확인 후 60K 본 학습 제출

smoke loss ~1.0 수준, step time, GPU mem 확인. 통과 시:

sbmr 5 "bash scripts/run_idm_train_robocasa_latent_A_v2_60k_8gpu.sh" --qos=core-extra --gres=gpu:8 -c 112 --mem 1600GB

60K 완료 후 24 ep eval

run_eval_latent_A_v2_24ep.sh 실행. v1 슬라이스 버전 (131.8mm / 5ep)과 비교: state 추가 supervision이 action mm error에 영향 주는지 검증. 예상: 동등하거나 미세 개선 (state supervision이 representation quality 향상).

State quality 평가

dump_v2 + 24-ep eval 완료 후: predicted_base_to_eef_pos vs GT mm error, quat angle error (deg), gripper open/close accuracy. v2가 DreamGen 파이프라인에 실질적으로 사용 가능한지 수치로 판단.

미해결: v1 ckpt 사용 불가 → 재학습 비용

[12:28] 영역이 v1에서 zero-supervision이었으므로 v1 ckpt 기반 fine-tune 효과 없음. 60K from-scratch가 유일 경로. 8 GPU × ~13h 예상.