#!/usr/bin/env bash
# 训练 ACT 策略 · 右臂取杯 · Jetson Orin Nano Super (8GB 统一内存)
#
# 环境选择说明（2026-08-20 实测）：
#   conda env `lerobot`  : torch 2.11.0+cu130  -> torch.cuda.is_available() == False
#                          （通用 cu130 轮子，在 Tegra 上用不了 GPU）
#   conda env `train310` : torch 2.8.0 Jetson 版 -> CUDA 可用，device=Orin, sm_87   ← 用这个
#
# 已修：该环境原装 numpy 2.2.6，而这版 torch 是按 NumPy 1.x 编译的，报
#       "Failed to initialize NumPy"。已降到 numpy 1.26.4。
#
# 内存：Jetson 的 CPU 与 GPU 共享同一块 7.4 GB。浏览器/编辑器占的每一 MB
#       都直接从训练可用显存里扣。开训前先关掉它们。
set -euo pipefail

PY=/home/robomates/miniconda3/envs/train310/bin/python
REPO_ID="${REPO_ID:-Suyang99/xlerobot-cup-grasp-20260820-0230}"
export HF_LEROBOT_HOME="${HF_LEROBOT_HOME:-/home/robomates/Robotic_challenge/data/huggingface/lerobot}"
export HF_HUB_OFFLINE="${HF_HUB_OFFLINE:-1}"     # 数据已在本地，不要联网重下

BATCH="${BATCH:-2}"          # Orin Nano 8GB 统一内存，2 是实测起步值
WORKERS="${WORKERS:-2}"      # 每个 worker 都要缓存解码后的帧，别调大
STEPS="${STEPS:-20000}"
SAVE_FREQ="${SAVE_FREQ:-2000}"
LOG_FREQ="${LOG_FREQ:-50}"
OUT="${OUT:-/home/robomates/Robotic_challenge/05-training/runs/act_$(date +%Y%m%d-%H%M)}"

cd /home/robomates/lerobot310
exec "$PY" -m lerobot.scripts.lerobot_train \
  --dataset.repo_id="$REPO_ID" \
  --policy.type=act \
  --policy.device=cuda \
  --policy.push_to_hub=false \
  --batch_size="$BATCH" \
  --num_workers="$WORKERS" \
  --steps="$STEPS" \
  --save_freq="$SAVE_FREQ" \
  --log_freq="$LOG_FREQ" \
  --output_dir="$OUT" \
  --job_name=act_cup_grasp_right \
  "$@"