Initial commit: DoVLA-CIL codebase (h=16 breakthrough) (part 2)

scripts/slurm/generate_6task_h16.sbatch

@@ -0,0 +1,87 @@
+#!/bin/bash
+#SBATCH --job-name=dovla_6task_h16
+#SBATCH --account=def-yalda_gpu
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=4
+#SBATCH --gres=gpu:nvidia_h100_80gb_hbm3_1g.10gb:1
+#SBATCH --mem=24G
+#SBATCH --time=08:00:00
+#SBATCH --output=outputs/hpc/logs/%x_%A_%a.out
+#SBATCH --error=outputs/hpc/logs/%x_%A_%a.err
+#SBATCH --array=0-5
+
+set -euo pipefail
+
+# Generate 6-task CIL collection with horizon=16 (vs baseline h=4)
+# Expected: oracle ceiling ~90%+ (vs 42.57% @ h=4)
+# This enables policy success 50-70%+ (vs 29.67% @ h=4)
+
+PROJECT_DIR="${PROJECT_DIR:-$SLURM_SUBMIT_DIR}"
+SCRATCH_ROOT="/scratch/$USER/dovla"
+SIF="$SCRATCH_ROOT/containers/pytorch_2.7.1_cuda12.8.sif"
+PYTHON="$SCRATCH_ROOT/envs/maniskill/bin/python"
+NATIVE_LIBS="$SCRATCH_ROOT/native_libs/lib"
+CPU_RENDER_LIBS="$SCRATCH_ROOT/cpu_render_libs"
+CA_BUNDLE="$SCRATCH_ROOT/ca-bundle.crt"
+VULKAN_ICD="$CPU_RENDER_LIBS/share/vulkan/icd.d/lvp_icd.x86_64.json"
+OUT_ROOT="${OUT_ROOT:-$SCRATCH_ROOT/experiments/six_task_h16_collection}"
+RUNTIME_DIR="/tmp/$USER/dovla-runtime-$SLURM_JOB_ID"
+CACHE_DIR="/tmp/$USER/dovla-mesa-$SLURM_JOB_ID"
+
+# Task array
+TASKS=(PickCube-v1 PushCube-v1 PullCube-v1 StackCube-v1 LiftPegUpright-v1 PegInsertionSide-v1)
+TASK=${TASKS[$SLURM_ARRAY_TASK_ID]}
+
+# Demo paths
+declare -A DEMOS
+DEMOS[PickCube-v1]="$SCRATCH_ROOT/maniskill_data/demos/PickCube-v1/rl/trajectory.none.pd_ee_delta_pose.physx_cuda.h5"
+DEMOS[PushCube-v1]="$SCRATCH_ROOT/maniskill_multitask_demos/PushCube-v1/rl/trajectory.none.pd_ee_delta_pose.physx_cuda.h5"
+DEMOS[PullCube-v1]="$SCRATCH_ROOT/maniskill_multitask_demos/PullCube-v1/rl/trajectory.none.pd_ee_delta_pose.physx_cuda.h5"
+DEMOS[StackCube-v1]="$SCRATCH_ROOT/maniskill_multitask_demos/StackCube-v1/rl/trajectory.none.pd_ee_delta_pose.physx_cuda.h5"
+DEMOS[LiftPegUpright-v1]="$SCRATCH_ROOT/maniskill_multitask_demos/LiftPegUpright-v1/rl/trajectory.none.pd_ee_delta_pose.physx_cuda.h5"
+DEMOS[PegInsertionSide-v1]="$SCRATCH_ROOT/maniskill_multitask_demos/PegInsertionSide-v1/rl/trajectory.h5"
+
+DEMO_PATH="${DEMOS[$TASK]}"
+OUT_DIR="$OUT_ROOT/$TASK"
+
+# Groups per task
+if [[ "$TASK" == "PickCube-v1" ]]; then
+  NUM_GROUPS=1000
+else
+  NUM_GROUPS=500
+fi
+
+module load StdEnv/2023 apptainer/1.4.5
+cd "$PROJECT_DIR"
+mkdir -p outputs/hpc/logs "$OUT_DIR" "$RUNTIME_DIR" "$CACHE_DIR"
+chmod 700 "$RUNTIME_DIR"
+
+export OMP_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 MKL_NUM_THREADS=1 LP_NUM_THREADS=1
+
+ENVS="LD_LIBRARY_PATH=$CPU_RENDER_LIBS/lib:$NATIVE_LIBS:/.singularity.d/libs,VK_ICD_FILENAMES=$VULKAN_ICD,VK_DRIVER_FILES=$VULKAN_ICD,XDG_RUNTIME_DIR=$RUNTIME_DIR,MESA_SHADER_CACHE_DIR=$CACHE_DIR,LIBGL_ALWAYS_SOFTWARE=1,LP_NUM_THREADS=1,SSL_CERT_FILE=$CA_BUNDLE,REQUESTS_CA_BUNDLE=$CA_BUNDLE,OMP_NUM_THREADS=1,OPENBLAS_NUM_THREADS=1,MKL_NUM_THREADS=1"
+
+echo "=================================================="
+echo "Task: $TASK"
+echo "Groups: $NUM_GROUPS"
+echo "Horizon: 16 (vs baseline 4)"
+echo "Demo: $DEMO_PATH"
+echo "Output: $OUT_DIR"
+echo "=================================================="
+
+apptainer exec --nv --env "$ENVS" \
+  "$SIF" "$PYTHON" scripts/generate_maniskill_lattice.py \
+  --demo "$DEMO_PATH" \
+  --out "$OUT_DIR" \
+  --env-id "$TASK" \
+  --num-groups "$NUM_GROUPS" \
+  --k 16 \
+  --horizon 16 \
+  --seed 0 \
+  --shard-size 1024 \
+  --sim-backend physx_cuda:0 \
+  --render-backend cpu \
+  --state-storage archive
+
+echo ""
+echo "✅ $TASK generation complete"

scripts/slurm/generate_cil_array.sbatch

@@ -0,0 +1,63 @@
+#!/bin/bash
+#SBATCH --job-name=${DOVLA_JOB_NAME:-dovla_cil_gen}
+#SBATCH --partition=${DOVLA_PARTITION:-compute}
+#SBATCH --array=${DOVLA_ARRAY:-0-9}
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=${DOVLA_CPUS_PER_TASK:-8}
+#SBATCH --gres=gpu:${DOVLA_GPUS_PER_TASK:-0}
+#SBATCH --mem=${DOVLA_MEM:-32G}
+#SBATCH --time=${DOVLA_TIME:-12:00:00}
+#SBATCH --output=${DOVLA_LOG_DIR:-logs/slurm}/%x_%A_%a.out
+#SBATCH --error=${DOVLA_LOG_DIR:-logs/slurm}/%x_%A_%a.err
+
+set -euo pipefail
+
+PROJECT_DIR="${PROJECT_DIR:-$PWD}"
+VENV_PATH="${VENV_PATH:-$PROJECT_DIR/.venv}"
+TASKS_PATH="${TASKS_PATH:-$PROJECT_DIR/data/tasks.jsonl}"
+OUT_ROOT="${OUT_ROOT:-$PROJECT_DIR/data/cil_array}"
+BACKEND="${BACKEND:-toy}"
+NUM_WORKERS="${NUM_WORKERS:-4}"
+STATES_PER_TASK="${STATES_PER_TASK:-1000}"
+K="${K:-32}"
+SHARD_SIZE="${SHARD_SIZE:-10000}"
+SEED_BASE="${SEED_BASE:-0}"
+RAY_ADDRESS="${RAY_ADDRESS:-}"
+RESUME_FLAG="${RESUME_FLAG:-}"
+
+mkdir -p "${DOVLA_LOG_DIR:-logs/slurm}" "$OUT_ROOT"
+cd "$PROJECT_DIR"
+
+if [ -f "$VENV_PATH/bin/activate" ]; then
+  # shellcheck disable=SC1091
+  source "$VENV_PATH/bin/activate"
+fi
+
+export OPENCLAUDE_BASE_URL="${OPENCLAUDE_BASE_URL:-https://open-claude.com/v1}"
+export OPENCLAUDE_MODEL="${OPENCLAUDE_MODEL:-<model>}"
+# Set OPENCLAUDE_API_KEY in the job environment or scheduler secret store. Do not echo it.
+
+SEED=$((SEED_BASE + SLURM_ARRAY_TASK_ID))
+OUT_DIR="$OUT_ROOT/part_${SLURM_ARRAY_TASK_ID}"
+
+CMD=(
+  python scripts/generate_cil_distributed.py
+  --backend "$BACKEND"
+  --tasks "$TASKS_PATH"
+  --out "$OUT_DIR"
+  --num-workers "$NUM_WORKERS"
+  --num-states-per-task "$STATES_PER_TASK"
+  --k "$K"
+  --seed "$SEED"
+  --shard-size "$SHARD_SIZE"
+)
+
+if [ -n "$RAY_ADDRESS" ]; then
+  CMD+=(--ray-address "$RAY_ADDRESS")
+fi
+if [ -n "$RESUME_FLAG" ]; then
+  CMD+=(--resume)
+fi
+
+"${CMD[@]}"

scripts/slurm/generate_embeddings.sbatch

@@ -0,0 +1,28 @@
+#!/bin/bash
+#SBATCH --job-name=gen_embeddings
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=8
+#SBATCH --mem=16000M
+#SBATCH --time=1:00:00
+#SBATCH --output=logs/gen_embeddings_%A.out
+#SBATCH --error=logs/gen_embeddings_%A.err
+
+set -euo pipefail
+
+PROJECT_DIR="${PROJECT_DIR:-$PWD}"
+cd "$PROJECT_DIR"
+
+source .venv/bin/activate
+
+echo "=== Generating Instruction Embeddings (Fast Parallel) ==="
+echo "Using 8 CPU cores for parallel encoding"
+echo ""
+
+python scripts/generate_instruction_embeddings.py \
+    --dataset /scratch/$USER/dovla/experiments/maniskill_presuccess_six_task_collection \
+    --output /scratch/$USER/dovla/experiments/instruction_embeddings.pkl \
+    --cache-dir /scratch/$USER/dovla/experiments/embedding_cache
+
+echo ""
+echo "✅ Embeddings generated successfully"

scripts/slurm/horizon_sweep_pickcube.sbatch

@@ -0,0 +1,88 @@
+#!/bin/bash
+#SBATCH --job-name=dovla_horizon_sweep
+#SBATCH --account=def-yalda_gpu
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=4
+#SBATCH --gres=gpu:nvidia_h100_80gb_hbm3_1g.10gb:1
+#SBATCH --mem=24G
+#SBATCH --time=01:30:00
+#SBATCH --output=outputs/hpc/logs/%x_%j.out
+#SBATCH --error=outputs/hpc/logs/%x_%j.err
+
+set -euo pipefail
+
+# DECISIVE EXPERIMENT: Does action horizon raise the oracle ceiling?
+# Generates PickCube CIL at horizon {4, 8, 16, 32}, measures oracle ceiling each.
+# Baseline (horizon=4) oracle for PickCube = 37.4%.
+
+PROJECT_DIR="${PROJECT_DIR:-$SLURM_SUBMIT_DIR}"
+SCRATCH_ROOT="/scratch/$USER/dovla"
+SIF="$SCRATCH_ROOT/containers/pytorch_2.7.1_cuda12.8.sif"
+PYTHON="$SCRATCH_ROOT/envs/maniskill/bin/python"
+NATIVE_LIBS="$SCRATCH_ROOT/native_libs/lib"
+CPU_RENDER_LIBS="$SCRATCH_ROOT/cpu_render_libs"
+CA_BUNDLE="$SCRATCH_ROOT/ca-bundle.crt"
+VULKAN_ICD="$CPU_RENDER_LIBS/share/vulkan/icd.d/lvp_icd.x86_64.json"
+DEMO="$SCRATCH_ROOT/maniskill_data/demos/PickCube-v1/rl/trajectory.none.pd_ee_delta_pose.physx_cuda.h5"
+OUT_ROOT="${OUT_ROOT:-$SCRATCH_ROOT/experiments/horizon_sweep_pickcube}"
+RUNTIME_DIR="/tmp/$USER/dovla-runtime-$SLURM_JOB_ID"
+CACHE_DIR="/tmp/$USER/dovla-mesa-$SLURM_JOB_ID"
+
+module load StdEnv/2023 apptainer/1.4.5
+cd "$PROJECT_DIR"
+mkdir -p outputs/hpc/logs "$OUT_ROOT" "$RUNTIME_DIR" "$CACHE_DIR"
+chmod 700 "$RUNTIME_DIR"
+
+export OMP_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 MKL_NUM_THREADS=1 LP_NUM_THREADS=1
+
+ENVS="LD_LIBRARY_PATH=$CPU_RENDER_LIBS/lib:$NATIVE_LIBS:/.singularity.d/libs,VK_ICD_FILENAMES=$VULKAN_ICD,VK_DRIVER_FILES=$VULKAN_ICD,XDG_RUNTIME_DIR=$RUNTIME_DIR,MESA_SHADER_CACHE_DIR=$CACHE_DIR,LIBGL_ALWAYS_SOFTWARE=1,LP_NUM_THREADS=1,SSL_CERT_FILE=$CA_BUNDLE,REQUESTS_CA_BUNDLE=$CA_BUNDLE,OMP_NUM_THREADS=1,OPENBLAS_NUM_THREADS=1,MKL_NUM_THREADS=1"
+
+for H in 4 8 16 32; do
+  OUT_DIR="$OUT_ROOT/h${H}"
+  echo "=================================================="
+  echo "Generating PickCube horizon=$H, 200 groups, K=16"
+  echo "=================================================="
+  apptainer exec --nv --env "$ENVS" \
+    "$SIF" "$PYTHON" scripts/generate_maniskill_lattice.py \
+    --demo "$DEMO" \
+    --out "$OUT_DIR" \
+    --env-id PickCube-v1 \
+    --num-groups 200 \
+    --k 16 \
+    --horizon "$H" \
+    --seed 0 \
+    --shard-size 1024 \
+    --sim-backend physx_cuda:0 \
+    --render-backend cpu \
+    --state-storage archive
+done
+
+echo ""
+echo "=================================================="
+echo "ORACLE CEILING BY HORIZON"
+echo "=================================================="
+apptainer exec --nv --env "$ENVS" "$SIF" "$PYTHON" - <<'PY'
+import sys; sys.path.insert(0,'.')
+from dovla_cil.data.datasets import CILDataset
+import os
+root=os.path.expandvars("/scratch/$USER/dovla/experiments/horizon_sweep_pickcube")
+print(f"{'horizon':>8} {'groups':>7} {'oracle':>8} {'expert':>8} {'mean_reward_spread':>18}")
+for H in [4,8,16,32]:
+    d=os.path.join(root,f"h{H}")
+    try:
+        ds=CILDataset(d)
+    except Exception as e:
+        print(f"{H:>8}  ERROR: {e}"); continue
+    n=len(ds.group_ids); orac=0; exp=0; spreads=[]
+    for gid in ds.group_ids:
+        recs=ds.get_group(gid)
+        if any(r.reward.terminal_success for r in recs): orac+=1
+        if any(r.candidate_type=='expert' and r.reward.terminal_success for r in recs): exp+=1
+        scores=[r.reward.score for r in recs]
+        spreads.append(max(scores)-min(scores))
+    ms=sum(spreads)/len(spreads) if spreads else 0
+    print(f"{H:>8} {n:>7} {orac/n:>8.4f} {exp/n:>8.4f} {ms:>18.4f}")
+print()
+print("Baseline reference: horizon=4 PickCube oracle in full collection = 0.3740")
+PY

scripts/slurm/install_smolvla_env.sbatch

@@ -0,0 +1,104 @@
+#!/bin/bash
+#SBATCH --job-name=dovla_smolvla_env
+#SBATCH --account=def-yalda_gpu
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=2
+#SBATCH --mem=8G
+#SBATCH --time=00:30:00
+#SBATCH --output=outputs/hpc/logs/%x_%j.out
+#SBATCH --error=outputs/hpc/logs/%x_%j.err
+
+set -euo pipefail
+
+PROJECT_DIR="${PROJECT_DIR:-$SLURM_SUBMIT_DIR}"
+SCRATCH_ROOT="${SCRATCH_ROOT:-/scratch/$USER/dovla}"
+CONTAINER="${CONTAINER:-$SCRATCH_ROOT/containers/pytorch_2.7.1_cuda12.8.sif}"
+ENV_DIR="${ENV_DIR:-$SCRATCH_ROOT/envs/smolvla}"
+LEROBOT_WHEEL="${LEROBOT_WHEEL:-$SCRATCH_ROOT/wheels/lerobot-0.4.3-py3-none-any.whl}"
+DRACCUS_WHEEL="${DRACCUS_WHEEL:-$SCRATCH_ROOT/wheels/draccus-0.10.0-py3-none-any.whl}"
+PYYAML_INCLUDE_WHEEL="${PYYAML_INCLUDE_WHEEL:-$SCRATCH_ROOT/wheels/pyyaml_include-1.4.1-py3-none-any.whl}"
+PYARROW_WHEEL="${PYARROW_WHEEL:-$SCRATCH_ROOT/wheels/pyarrow-17.0.0-cp311-cp311-linux_x86_64.whl}"
+DATASETS_WHEEL="${DATASETS_WHEEL:-/cvmfs/soft.computecanada.ca/custom/python/wheelhouse/generic/datasets-4.0.0+computecanada-py3-none-any.whl}"
+WHEELHOUSE_ARCH="${WHEELHOUSE_ARCH:-/cvmfs/soft.computecanada.ca/custom/python/wheelhouse/gentoo2023/x86-64-v3}"
+WHEELHOUSE_GENERIC="${WHEELHOUSE_GENERIC:-/cvmfs/soft.computecanada.ca/custom/python/wheelhouse/gentoo2023/generic}"
+
+cd "$PROJECT_DIR"
+mkdir -p outputs/hpc/logs "$SCRATCH_ROOT/envs"
+module load StdEnv/2023 apptainer/1.4.5
+
+for WHEEL in \
+  "$LEROBOT_WHEEL" \
+  "$DRACCUS_WHEEL" \
+  "$PYYAML_INCLUDE_WHEEL" \
+  "$PYARROW_WHEEL" \
+  "$DATASETS_WHEEL"; do
+  if [[ ! -f "$WHEEL" ]]; then
+    echo "Missing pinned runtime wheel: $WHEEL" >&2
+    echo "Stage all pinned wheels before submitting this offline job." >&2
+    exit 2
+  fi
+done
+
+if [[ ! -x "$ENV_DIR/bin/python" ]]; then
+  apptainer exec \
+    -B "$SCRATCH_ROOT:$SCRATCH_ROOT" \
+    "$CONTAINER" \
+    /opt/conda/bin/python -m venv --system-site-packages "$ENV_DIR"
+fi
+
+apptainer exec \
+  -B "$SCRATCH_ROOT:$SCRATCH_ROOT" \
+  -B "$PROJECT_DIR:$PROJECT_DIR" \
+  -B /cvmfs:/cvmfs \
+  "$CONTAINER" \
+  "$ENV_DIR/bin/python" -c \
+  "from itertools import islice; from packaging.tags import sys_tags; print('supported_tags', [str(tag) for tag in islice(sys_tags(), 12)])"
+
+apptainer exec \
+  -B "$SCRATCH_ROOT:$SCRATCH_ROOT" \
+  -B "$PROJECT_DIR:$PROJECT_DIR" \
+  -B /cvmfs:/cvmfs \
+  "$CONTAINER" \
+  "$ENV_DIR/bin/python" -m pip install \
+    --no-index \
+    --find-links "$WHEELHOUSE_ARCH" \
+    --find-links "$WHEELHOUSE_GENERIC" \
+    "transformers==4.57.6+computecanada" \
+    "huggingface-hub==0.35.3+computecanada" \
+    "accelerate==1.10.1+computecanada" \
+    "num2words==0.5.14+computecanada" \
+    "typing-inspect==0.9.0+computecanada" \
+    "mergedeep==1.3.4+computecanada" \
+    "toml==0.10.2+computecanada" \
+    "einops==0.8.1+computecanada" \
+    "dill==0.3.8+computecanada" \
+    "multiprocess==0.70.16+computecanada" \
+    "xxhash==3.5.0+computecanada" \
+    "pandas==2.2.3+computecanada" \
+    "fsspec==2025.3.0+computecanada" \
+    "setuptools==80.9.0+computecanada" \
+    "imageio==2.37.0+computecanada" \
+    "imageio-ffmpeg==0.6.0+computecanada"
+
+apptainer exec \
+  -B "$SCRATCH_ROOT:$SCRATCH_ROOT" \
+  -B "$PROJECT_DIR:$PROJECT_DIR" \
+  -B /cvmfs:/cvmfs \
+  "$CONTAINER" \
+  "$ENV_DIR/bin/python" -m pip install \
+    --no-index \
+    --no-deps \
+    "$PYYAML_INCLUDE_WHEEL" \
+    "$PYARROW_WHEEL" \
+    "$DATASETS_WHEEL" \
+    "$DRACCUS_WHEEL" \
+    "$LEROBOT_WHEEL"
+
+apptainer exec \
+  -B "$SCRATCH_ROOT:$SCRATCH_ROOT" \
+  -B "$PROJECT_DIR:$PROJECT_DIR" \
+  --env "PYTHONPATH=$PROJECT_DIR" \
+  "$CONTAINER" \
+  "$ENV_DIR/bin/python" -c \
+  "import accelerate, datasets, importlib.util, lerobot, pyarrow, transformers; from dovla_cil.eval.smolvla_runtime import import_smolvla_classes; SmolVLAPolicy, SmolVLAConfig = import_smolvla_classes(); print('lerobot', lerobot.__version__); print('transformers', transformers.__version__); print('accelerate', accelerate.__version__); print('datasets', datasets.__version__); print('pyarrow', pyarrow.__version__); print('policy_import', SmolVLAPolicy.__name__, SmolVLAConfig.__name__); [print(name, bool(importlib.util.find_spec(name))) for name in ('draccus', 'typing_inspect', 'gymnasium', 'einops', 'safetensors')]"

scripts/slurm/make_maniskill_collection.sbatch

@@ -0,0 +1,32 @@
+#!/bin/bash
+#SBATCH --job-name=dovla_ms_collect
+#SBATCH --account=def-yalda_cpu
+#SBATCH --partition=cpubase_bycore_b1
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=2
+#SBATCH --mem=8G
+#SBATCH --time=00:20:00
+#SBATCH --output=outputs/hpc/logs/%x_%j.out
+#SBATCH --error=outputs/hpc/logs/%x_%j.err
+
+set -euo pipefail
+
+PROJECT_DIR="${PROJECT_DIR:-$SLURM_SUBMIT_DIR}"
+PICKCUBE_DATA="${PICKCUBE_DATA:?Set PICKCUBE_DATA}"
+MULTITASK_ROOT="${MULTITASK_ROOT:?Set MULTITASK_ROOT}"
+COLLECTION_OUT="${COLLECTION_OUT:?Set COLLECTION_OUT}"
+COLLECTION_NAME="${COLLECTION_NAME:-maniskill-six-task-k16}"
+PYTHON="${PYTHON:-$PROJECT_DIR/.venv/bin/python}"
+
+cd "$PROJECT_DIR"
+"$PYTHON" scripts/make_cil_collection.py \
+  --name "$COLLECTION_NAME" \
+  --out "$COLLECTION_OUT" \
+  --sources \
+    "$PICKCUBE_DATA" \
+    "$MULTITASK_ROOT/PushCube-v1" \
+    "$MULTITASK_ROOT/PullCube-v1" \
+    "$MULTITASK_ROOT/StackCube-v1" \
+    "$MULTITASK_ROOT/LiftPegUpright-v1" \
+    "$MULTITASK_ROOT/PegInsertionSide-v1"

scripts/slurm/maniskill_lattice_debug.sbatch

@@ -0,0 +1,111 @@
+#!/bin/bash
+#SBATCH --job-name=dovla_ms_debug
+#SBATCH --account=def-yalda_gpu
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=4
+# Physics uses CUDA; state-mode material creation uses the CPU Vulkan renderer to avoid
+# Vulkan/CUDA device-ordinal mismatches on shared four-GPU nodes.
+#SBATCH --gres=gpu:nvidia_h100_80gb_hbm3_1g.10gb:1
+#SBATCH --mem=24G
+#SBATCH --time=00:20:00
+#SBATCH --output=outputs/hpc/logs/%x_%j.out
+#SBATCH --error=outputs/hpc/logs/%x_%j.err
+
+set -euo pipefail
+
+PROJECT_DIR="${PROJECT_DIR:-$SLURM_SUBMIT_DIR}"
+SCRATCH_ROOT="/scratch/$USER/dovla"
+SIF="$SCRATCH_ROOT/containers/pytorch_2.7.1_cuda12.8.sif"
+PYTHON="$SCRATCH_ROOT/envs/maniskill/bin/python"
+NATIVE_LIBS="$SCRATCH_ROOT/native_libs/lib"
+CPU_RENDER_LIBS="$SCRATCH_ROOT/cpu_render_libs"
+CA_BUNDLE="$SCRATCH_ROOT/ca-bundle.crt"
+VULKAN_ICD="$CPU_RENDER_LIBS/share/vulkan/icd.d/lvp_icd.x86_64.json"
+DEMO="$SCRATCH_ROOT/maniskill_data/demos/PickCube-v1/rl/trajectory.none.pd_ee_delta_pose.physx_cuda.h5"
+OUT_DIR="${OUT_DIR:-$PROJECT_DIR/outputs/hpc/maniskill_debug_cil}"
+RUNTIME_DIR="/tmp/$USER/dovla-runtime-$SLURM_JOB_ID"
+CACHE_DIR="/tmp/$USER/dovla-mesa-$SLURM_JOB_ID"
+
+module load StdEnv/2023 apptainer/1.4.5
+cd "$PROJECT_DIR"
+mkdir -p outputs/hpc/logs "$OUT_DIR" "$RUNTIME_DIR" "$CACHE_DIR"
+chmod 700 "$RUNTIME_DIR"
+
+export OMP_NUM_THREADS=1
+export OPENBLAS_NUM_THREADS=1
+export MKL_NUM_THREADS=1
+export LP_NUM_THREADS=1
+
+apptainer exec --nv \
+  --env "LD_LIBRARY_PATH=$CPU_RENDER_LIBS/lib:$NATIVE_LIBS:/.singularity.d/libs,VK_ICD_FILENAMES=$VULKAN_ICD,VK_DRIVER_FILES=$VULKAN_ICD,XDG_RUNTIME_DIR=$RUNTIME_DIR,MESA_SHADER_CACHE_DIR=$CACHE_DIR,LIBGL_ALWAYS_SOFTWARE=1,LP_NUM_THREADS=1,SSL_CERT_FILE=$CA_BUNDLE,REQUESTS_CA_BUNDLE=$CA_BUNDLE,OMP_NUM_THREADS=1,OPENBLAS_NUM_THREADS=1,MKL_NUM_THREADS=1" \
+  "$SIF" "$PYTHON" - <<'PY'
+import gymnasium as gym
+import mani_skill
+import os
+import torch
+
+print("torch", torch.__version__, "cuda", torch.cuda.is_available(), torch.cuda.get_device_name(0))
+print("vulkan_icd", os.environ.get("VK_ICD_FILENAMES"), "cuda_visible", os.environ.get("CUDA_VISIBLE_DEVICES"))
+env = gym.make(
+    "PickCube-v1",
+    num_envs=1,
+    obs_mode="state",
+    control_mode="pd_ee_delta_pose",
+    render_mode=None,
+    sim_backend="physx_cuda:0",
+    render_backend="cpu",
+)
+env.reset(seed=7)
+state = {
+    section: {name: value.clone() for name, value in values.items()}
+    for section, values in env.unwrapped.get_state_dict().items()
+}
+action = torch.zeros((1, 7), dtype=torch.float32, device=env.unwrapped.device)
+env.unwrapped.set_state_dict(state)
+env.unwrapped.agent.controller.reset()
+restored = env.unwrapped.get_state_dict()
+max_error = max(
+    float(torch.max(torch.abs(state[section][name] - restored[section][name])).cpu())
+    for section in state
+    for name in state[section]
+)
+print("state_restore_max_error", max_error)
+assert max_error <= 1e-6
+
+env.unwrapped.step(action)
+next_state_1 = {
+    section: {name: value.clone() for name, value in values.items()}
+    for section, values in env.unwrapped.get_state_dict().items()
+}
+env.unwrapped.set_state_dict(state)
+env.unwrapped.agent.controller.reset()
+env.unwrapped.step(action)
+next_state_2 = env.unwrapped.get_state_dict()
+branch_error = max(
+    float(torch.max(torch.abs(next_state_1[section][name] - next_state_2[section][name])).cpu())
+    for section in next_state_1
+    for name in next_state_1[section]
+)
+print("deterministic_branch_max_error", branch_error)
+assert branch_error <= 1e-5
+env.close()
+PY
+
+apptainer exec --nv \
+  --env "LD_LIBRARY_PATH=$CPU_RENDER_LIBS/lib:$NATIVE_LIBS:/.singularity.d/libs,VK_ICD_FILENAMES=$VULKAN_ICD,VK_DRIVER_FILES=$VULKAN_ICD,XDG_RUNTIME_DIR=$RUNTIME_DIR,MESA_SHADER_CACHE_DIR=$CACHE_DIR,LIBGL_ALWAYS_SOFTWARE=1,LP_NUM_THREADS=1,SSL_CERT_FILE=$CA_BUNDLE,REQUESTS_CA_BUNDLE=$CA_BUNDLE,OMP_NUM_THREADS=1,OPENBLAS_NUM_THREADS=1,MKL_NUM_THREADS=1" \
+  "$SIF" "$PYTHON" scripts/generate_maniskill_lattice.py \
+  --demo "$DEMO" \
+  --out "$OUT_DIR" \
+  --num-groups 8 \
+  --k 4 \
+  --horizon 4 \
+  --seed 0 \
+  --shard-size 32 \
+  --sim-backend physx_cuda:0 \
+  --render-backend cpu \
+  --state-storage archive
+
+apptainer exec --nv \
+  --env "LD_LIBRARY_PATH=$CPU_RENDER_LIBS/lib:$NATIVE_LIBS:/.singularity.d/libs,VK_ICD_FILENAMES=$VULKAN_ICD,VK_DRIVER_FILES=$VULKAN_ICD,XDG_RUNTIME_DIR=$RUNTIME_DIR,MESA_SHADER_CACHE_DIR=$CACHE_DIR,LIBGL_ALWAYS_SOFTWARE=1,LP_NUM_THREADS=1,SSL_CERT_FILE=$CA_BUNDLE,REQUESTS_CA_BUNDLE=$CA_BUNDLE,OMP_NUM_THREADS=1,OPENBLAS_NUM_THREADS=1,MKL_NUM_THREADS=1" \
+  "$SIF" "$PYTHON" scripts/inspect_shard.py "$OUT_DIR/manifest.json"

scripts/slurm/maniskill_lattice_full.sbatch

@@ -0,0 +1,80 @@
+#!/bin/bash
+#SBATCH --job-name=dovla_ms_full
+#SBATCH --account=def-yalda_gpu
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=8
+#SBATCH --gres=gpu:nvidia_h100_80gb_hbm3_1g.10gb:1
+#SBATCH --mem=32G
+#SBATCH --time=02:00:00
+#SBATCH --output=outputs/hpc/logs/%x_%j.out
+#SBATCH --error=outputs/hpc/logs/%x_%j.err
+
+set -euo pipefail
+
+PROJECT_DIR="${PROJECT_DIR:-$SLURM_SUBMIT_DIR}"
+SCRATCH_ROOT="/scratch/$USER/dovla"
+SIF="$SCRATCH_ROOT/containers/pytorch_2.7.1_cuda12.8.sif"
+PYTHON="$SCRATCH_ROOT/envs/maniskill/bin/python"
+NATIVE_LIBS="$SCRATCH_ROOT/native_libs/lib"
+CPU_RENDER_LIBS="$SCRATCH_ROOT/cpu_render_libs"
+CA_BUNDLE="$SCRATCH_ROOT/ca-bundle.crt"
+VULKAN_ICD="$CPU_RENDER_LIBS/share/vulkan/icd.d/lvp_icd.x86_64.json"
+DEMO="${DEMO:-$SCRATCH_ROOT/maniskill_data/demos/PickCube-v1/rl/trajectory.none.pd_ee_delta_pose.physx_cuda.h5}"
+ENV_ID="${ENV_ID:-PickCube-v1}"
+CONTROL_MODE="${CONTROL_MODE:-pd_ee_delta_pose}"
+
+NUM_GROUPS="${NUM_GROUPS:-1000}"
+GROUP_OFFSET="${GROUP_OFFSET:-0}"
+K="${K:-16}"
+HORIZON="${HORIZON:-4}"
+SEED="${SEED:-0}"
+SHARD_SIZE="${SHARD_SIZE:-2048}"
+STATE_BATCH_SIZE="${STATE_BATCH_SIZE:-16}"
+OBS_MODE="${OBS_MODE:-state}"
+IMAGE_QUALITY="${IMAGE_QUALITY:-90}"
+CANDIDATE_MODE="${CANDIDATE_MODE:-structured}"
+OUT_DIR="${OUT_DIR:-$PROJECT_DIR/outputs/hpc/maniskill_full_k${K}_n${NUM_GROUPS}_seed${SEED}}"
+RUNTIME_DIR="/tmp/$USER/dovla-runtime-$SLURM_JOB_ID"
+CACHE_DIR="/tmp/$USER/dovla-mesa-$SLURM_JOB_ID"
+
+module load StdEnv/2023 apptainer/1.4.5
+cd "$PROJECT_DIR"
+mkdir -p outputs/hpc/logs "$OUT_DIR" "$RUNTIME_DIR" "$CACHE_DIR"
+chmod 700 "$RUNTIME_DIR"
+
+export OMP_NUM_THREADS=1
+export OPENBLAS_NUM_THREADS=1
+export MKL_NUM_THREADS=1
+export LP_NUM_THREADS=1
+
+if [[ -f "$OUT_DIR/manifest.json" ]]; then
+  echo "completed manifest already exists: $OUT_DIR/manifest.json"
+  exit 0
+fi
+
+apptainer exec --nv \
+  --env "LD_LIBRARY_PATH=$CPU_RENDER_LIBS/lib:$NATIVE_LIBS:/.singularity.d/libs,VK_ICD_FILENAMES=$VULKAN_ICD,VK_DRIVER_FILES=$VULKAN_ICD,XDG_RUNTIME_DIR=$RUNTIME_DIR,MESA_SHADER_CACHE_DIR=$CACHE_DIR,LIBGL_ALWAYS_SOFTWARE=1,LP_NUM_THREADS=1,SSL_CERT_FILE=$CA_BUNDLE,REQUESTS_CA_BUNDLE=$CA_BUNDLE,OMP_NUM_THREADS=1,OPENBLAS_NUM_THREADS=1,MKL_NUM_THREADS=1" \
+  "$SIF" "$PYTHON" scripts/generate_maniskill_lattice.py \
+  --demo "$DEMO" \
+  --env-id "$ENV_ID" \
+  --control-mode "$CONTROL_MODE" \
+  --out "$OUT_DIR" \
+  --num-groups "$NUM_GROUPS" \
+  --group-offset "$GROUP_OFFSET" \
+  --k "$K" \
+  --horizon "$HORIZON" \
+  --seed "$SEED" \
+  --shard-size "$SHARD_SIZE" \
+  --obs-mode "$OBS_MODE" \
+  --image-quality "$IMAGE_QUALITY" \
+  --sim-backend physx_cuda:0 \
+  --render-backend cpu \
+  --parallel-branches \
+  --state-batch-size "$STATE_BATCH_SIZE" \
+  --state-storage archive \
+  --candidate-mode "$CANDIDATE_MODE"
+
+apptainer exec \
+  --env "OMP_NUM_THREADS=1,OPENBLAS_NUM_THREADS=1,MKL_NUM_THREADS=1" \
+  "$SIF" "$PYTHON" scripts/inspect_shard.py "$OUT_DIR/manifest.json"

scripts/slurm/maniskill_multitask_pilot.sbatch

@@ -0,0 +1,63 @@
+#!/bin/bash
+#SBATCH --job-name=dovla_ms_multi
+#SBATCH --account=def-yalda_gpu
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=8
+#SBATCH --gres=gpu:nvidia_h100_80gb_hbm3_1g.10gb:1
+#SBATCH --mem=32G
+#SBATCH --time=00:20:00
+#SBATCH --array=0-4%2
+#SBATCH --output=outputs/hpc/logs/%x_%A_%a.out
+#SBATCH --error=outputs/hpc/logs/%x_%A_%a.err
+
+set -euo pipefail
+
+PROJECT_DIR="${PROJECT_DIR:-$SLURM_SUBMIT_DIR}"
+SCRATCH_ROOT="/scratch/$USER/dovla"
+DEMO_ROOT="$SCRATCH_ROOT/maniskill_multitask_demos"
+MULTITASK_OUT_ROOT="${MULTITASK_OUT_ROOT:-$SCRATCH_ROOT/experiments/maniskill_multitask_pilot}"
+
+case "${SLURM_ARRAY_TASK_ID:-0}" in
+  0)
+    ENV_ID="PushCube-v1"
+    CONTROL_MODE="pd_ee_delta_pose"
+    DEMO="$DEMO_ROOT/PushCube-v1/rl/trajectory.none.pd_ee_delta_pose.physx_cuda.h5"
+    ;;
+  1)
+    ENV_ID="PullCube-v1"
+    CONTROL_MODE="pd_ee_delta_pose"
+    DEMO="$DEMO_ROOT/PullCube-v1/rl/trajectory.none.pd_ee_delta_pose.physx_cuda.h5"
+    ;;
+  2)
+    ENV_ID="StackCube-v1"
+    CONTROL_MODE="pd_ee_delta_pose"
+    DEMO="$DEMO_ROOT/StackCube-v1/rl/trajectory.none.pd_ee_delta_pose.physx_cuda.h5"
+    ;;
+  3)
+    ENV_ID="LiftPegUpright-v1"
+    CONTROL_MODE="pd_ee_delta_pose"
+    DEMO="$DEMO_ROOT/LiftPegUpright-v1/rl/trajectory.none.pd_ee_delta_pose.physx_cuda.h5"
+    ;;
+  4)
+    ENV_ID="PegInsertionSide-v1"
+    CONTROL_MODE="pd_joint_pos"
+    DEMO="$DEMO_ROOT/PegInsertionSide-v1/motionplanning/trajectory.h5"
+    ;;
+  *)
+    echo "unsupported array index" >&2
+    exit 2
+    ;;
+esac
+
+export PROJECT_DIR DEMO ENV_ID CONTROL_MODE
+export NUM_GROUPS="${NUM_GROUPS:-16}"
+export K="${K:-8}"
+export HORIZON="${HORIZON:-4}"
+export STATE_BATCH_SIZE="${STATE_BATCH_SIZE:-8}"
+export OBS_MODE=state
+export SHARD_SIZE="${SHARD_SIZE:-256}"
+export STATE_STORAGE=archive
+export OUT_DIR="$MULTITASK_OUT_ROOT/$ENV_ID"
+
+exec bash "$PROJECT_DIR/scripts/slurm/maniskill_lattice_full.sbatch"

scripts/slurm/phase_a1_generate_10k.sbatch

@@ -0,0 +1,93 @@
+#!/bin/bash
+#SBATCH --job-name=dovla_10k_gen
+#SBATCH --partition=${DOVLA_PARTITION:-compute}
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=16
+#SBATCH --gres=gpu:1
+#SBATCH --mem=64G
+#SBATCH --time=48:00:00
+#SBATCH --output=logs/phase_a_10k_gen_%j.out
+#SBATCH --error=logs/phase_a_10k_gen_%j.err
+
+set -euo pipefail
+
+# Phase A1: Generate 10K groups dataset for performance improvement
+# This scales from current 3,500 to 10,000 groups
+# Expected: +5-10% success improvement
+
+PROJECT_DIR="${PROJECT_DIR:-$PWD}"
+cd "$PROJECT_DIR"
+
+# Activate environment
+if [ -f ".venv/bin/activate" ]; then
+    source .venv/bin/activate
+fi
+
+# Configuration
+DEMO_DIR="/scratch/$USER/dovla/demonstrations/maniskill"
+OUT_DIR="/scratch/$USER/dovla/experiments/phase_a_10k_collection"
+K=16
+STATE_BATCH_SIZE=16
+
+# Task configuration: 6 tasks with more groups each
+declare -A TASK_GROUPS=(
+    ["PickCube-v1"]=2000      # Increase from 1000
+    ["PushCube-v1"]=2000      # Increase from 500
+    ["PullCube-v1"]=1500      # Increase from 500
+    ["StackCube-v1"]=1500     # Increase from 500
+    ["LiftPegUpright-v1"]=1500 # Increase from 500
+    ["PegInsertionSide-v1"]=1500 # Increase from 500
+)
+
+mkdir -p "$OUT_DIR" logs
+
+echo "=== Phase A1: Generating 10K Group Collection ==="
+echo "Target: 10,000 groups, 160,000 records (K=$K)"
+echo "Expected improvement: +5-10% policy success"
+echo ""
+
+for TASK in "${!TASK_GROUPS[@]}"; do
+    NUM_GROUPS="${TASK_GROUPS[$TASK]}"
+    DEMO_FILE="$DEMO_DIR/${TASK}.h5"
+    TASK_OUT="$OUT_DIR/${TASK}_k${K}_n${NUM_GROUPS}"
+
+    if [ ! -f "$DEMO_FILE" ]; then
+        echo "⚠️  Demo file not found: $DEMO_FILE"
+        echo "   Skipping $TASK"
+        continue
+    fi
+
+    echo "Generating $TASK: $NUM_GROUPS groups..."
+
+    python scripts/generate_maniskill_lattice.py \
+        --demo "$DEMO_FILE" \
+        --env-id "$TASK" \
+        --control-mode pd_ee_delta_pose \
+        --out "$TASK_OUT" \
+        --num-groups "$NUM_GROUPS" \
+        --k "$K" \
+        --state-batch-size "$STATE_BATCH_SIZE" \
+        --seed 42 \
+        --pre-success-only
+
+    if [ $? -eq 0 ]; then
+        echo "✅ $TASK complete: $NUM_GROUPS groups"
+    else
+        echo "❌ $TASK failed"
+        exit 1
+    fi
+    echo ""
+done
+
+echo "=== Merging into unified collection ==="
+
+python scripts/make_cil_collection.py \
+    --source-dirs "$OUT_DIR"/*/  \
+    --out "$OUT_DIR/merged_10k" \
+    --name "phase_a_10k_collection"
+
+echo "✅ Phase A1 complete: 10K group collection ready"
+echo "   Location: $OUT_DIR/merged_10k"
+echo ""
+echo "Next: Run phase_a2_train_large_model.sbatch"

scripts/slurm/phase_a1_generate_10k_enhanced.sbatch

@@ -0,0 +1,120 @@
+#!/bin/bash
+#SBATCH --job-name=dovla_10k_gen
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=16
+#SBATCH --gres=gpu:1
+#SBATCH --mem=64000M
+#SBATCH --time=96:00:00
+#SBATCH --output=logs/phase_a1_10k_gen_%j.out
+#SBATCH --error=logs/phase_a1_10k_gen_%j.err
+
+set -euo pipefail
+
+# Phase A1: Enhanced 10K Generation
+# Target: 50%+ policy success with optimizations
+
+PROJECT_DIR="${PROJECT_DIR:-$PWD}"
+cd "$PROJECT_DIR"
+
+source .venv/bin/activate
+
+OUT_DIR="/scratch/$USER/dovla/experiments/phase_a1_10k_collection"
+K=16
+STATE_BATCH_SIZE=16
+
+echo "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "="
+echo "Phase A1: Enhanced 10K Generation for 50%+ Target"
+echo "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "="
+echo ""
+echo "Strategy:"
+echo "  - 10,000 groups (vs 3,500 current)"
+echo "  - 160,000 records total"
+echo "  - K=16 interventions per group"
+echo "  - Optimized for diverse counterfactuals"
+echo ""
+echo "Expected outcome: 42-50% policy success"
+echo ""
+
+# Task distribution (balanced across difficulty)
+declare -A TASK_GROUPS=(
+    ["PickCube-v1"]=1800       # Easy
+    ["PushCube-v1"]=1800       # Easy
+    ["PullCube-v1"]=1600       # Medium
+    ["StackCube-v1"]=1600      # Medium-Hard
+    ["LiftPegUpright-v1"]=1600 # Medium-Hard
+    ["PegInsertionSide-v1"]=1600 # Hard
+)
+
+TOTAL_GROUPS=0
+for count in "${TASK_GROUPS[@]}"; do
+    TOTAL_GROUPS=$((TOTAL_GROUPS + count))
+done
+
+echo "Task distribution (total: $TOTAL_GROUPS groups):"
+for TASK in "${!TASK_GROUPS[@]}"; do
+    echo "  ${TASK}: ${TASK_GROUPS[$TASK]} groups"
+done
+echo ""
+
+# Generate each task
+for TASK in "${!TASK_GROUPS[@]}"; do
+    NUM_GROUPS="${TASK_GROUPS[$TASK]}"
+
+    TASK_OUT="$OUT_DIR/${TASK}_k${K}_n${NUM_GROUPS}"
+
+    if [ -d "$TASK_OUT/merged" ]; then
+        echo "✓ $TASK already generated, skipping"
+        continue
+    fi
+
+    echo "Generating $TASK: $NUM_GROUPS groups..."
+    echo "  Start: $(date)"
+
+    # Determine demo path
+    DEMO_PATH="/scratch/$USER/dovla/demos/maniskill/${TASK%.v1}.h5"
+    if [ ! -f "$DEMO_PATH" ]; then
+        echo "  ⚠️ Demo not found at $DEMO_PATH, trying alternate location..."
+        DEMO_PATH="/scratch/$USER/dovla/experiments/maniskill_presuccess_six_task_collection/${TASK}/demos/demo.h5"
+    fi
+
+    if [ ! -f "$DEMO_PATH" ]; then
+        echo "  ❌ Demo not found, skipping $TASK"
+        continue
+    fi
+
+    python scripts/generate_maniskill_lattice.py \
+        --demo "$DEMO_PATH" \
+        --env-id "$TASK" \
+        --control-mode pd_ee_delta_pose \
+        --out "$TASK_OUT" \
+        --num-groups "$NUM_GROUPS" \
+        --k "$K" \
+        --state-batch-size "$STATE_BATCH_SIZE" \
+        --seed 42
+
+    echo "  ✅ Complete: $(date)"
+    echo ""
+done
+
+echo "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "="
+echo "Merging all tasks into unified collection"
+echo "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "="
+
+python scripts/make_cil_collection.py \
+    --source-dirs "$OUT_DIR"/*/merged \
+    --out "$OUT_DIR/merged_10k" \
+    --name "phase_a1_10k_enhanced"
+
+echo ""
+echo "✅ Phase A1 Enhanced Generation Complete!"
+echo ""
+echo "Output: $OUT_DIR/merged_10k"
+echo "Total groups: $TOTAL_GROUPS"
+echo "Total records: $((TOTAL_GROUPS * K))"
+echo ""
+echo "Next: Train enhanced model with:"
+echo "  - Hidden dim: 512"
+echo "  - Epochs: 150"
+echo "  - LR: 0.0003"
+echo "  - Enhanced loss weights"

scripts/slurm/phase_a1_revised_enhanced.sbatch

@@ -0,0 +1,65 @@
+#!/bin/bash
+#SBATCH --job-name=dovla_enhanced_train
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=8
+#SBATCH --gres=gpu:1
+#SBATCH --mem=64000M
+#SBATCH --time=48:00:00
+#SBATCH --output=logs/phase_a1_enhanced_single_%A_%a.out
+#SBATCH --error=logs/phase_a1_enhanced_single_%A_%a.err
+#SBATCH --array=0-2
+
+set -euo pipefail
+
+# Phase A1-Revised: Enhanced Training on Existing 3.5K Data
+# Target: 45%+ with better training, no new data needed
+
+PROJECT_DIR="${PROJECT_DIR:-$PWD}"
+cd "$PROJECT_DIR"
+
+source .venv/bin/activate
+
+DATASET="/scratch/$USER/dovla/experiments/maniskill_presuccess_six_task_collection"
+OUT_DIR="/scratch/$USER/dovla/experiments/phase_a1_revised_enhanced"
+SEED=$SLURM_ARRAY_TASK_ID
+
+mkdir -p "$OUT_DIR/seed_$SEED" logs
+
+echo "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "="
+echo "Phase A1-Revised: Enhanced Training (Existing Data)"
+echo "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "="
+echo ""
+echo "Strategy: Better training, not more data"
+echo "Seed: $SEED"
+echo "Dataset: 3,500 groups (existing)"
+echo "Model: h=256 (best from Phase A4)"
+echo "Training: 200 epochs with cosine schedule"
+echo ""
+echo "Target: 45%+ policy success"
+echo ""
+
+python scripts/train_dovla.py \
+    --dataset "$DATASET" \
+    --out "$OUT_DIR/seed_$SEED" \
+    --objective lattice_field \
+    --hidden-dim 256 \
+    --action-horizon 4 \
+    --epochs 200 \
+    --batch-groups 16 \
+    --records-per-group 8 \
+    --lr 0.0003 \
+    --weight-decay 0.01 \
+    --device auto \
+    --seed $SEED \
+    --observation-mode state \
+    --loss-weight bc=1.0 \
+    --loss-weight field_effect=1.5 \
+    --loss-weight field_potential=1.0 \
+    --loss-weight field_preference=0.8 \
+    --loss-weight field_anchor=0.2
+
+echo ""
+echo "✅ Phase A1-Revised enhanced training complete (seed $SEED)"
+echo ""
+echo "Next: Evaluate and check if 45%+ achieved"

scripts/slurm/phase_a1b_train_enhanced.sbatch

@@ -0,0 +1,63 @@
+#!/bin/bash
+#SBATCH --job-name=dovla_enhanced_train
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=8
+#SBATCH --gres=gpu:1
+#SBATCH --mem=64000M
+#SBATCH --time=120:00:00
+#SBATCH --output=logs/phase_a1b_enhanced_train_%A_%a.out
+#SBATCH --error=logs/phase_a1b_enhanced_train_%A_%a.err
+#SBATCH --array=0-2
+
+set -euo pipefail
+
+# Phase A1b: Enhanced Training for 50%+ Target
+
+PROJECT_DIR="${PROJECT_DIR:-$PWD}"
+cd "$PROJECT_DIR"
+
+source .venv/bin/activate
+
+DATASET="/scratch/$USER/dovla/experiments/phase_a1_10k_collection/merged_10k"
+OUT_DIR="/scratch/$USER/dovla/experiments/phase_a1b_enhanced_model"
+SEED=$SLURM_ARRAY_TASK_ID
+
+mkdir -p "$OUT_DIR/seed_$SEED" logs
+
+echo "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "="
+echo "Phase A1b: Enhanced Training for 50%+ Target"
+echo "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "="
+echo ""
+echo "Seed: $SEED"
+echo "Dataset: 10,000 groups, 160,000 records"
+echo "Model: hidden_dim=512 (optimal size for 10K data)"
+echo "Training: 150 epochs with warmup + decay"
+echo ""
+echo "Target: 50%+ policy success"
+echo ""
+
+python scripts/train_dovla.py \
+    --dataset "$DATASET" \
+    --out "$OUT_DIR/seed_$SEED" \
+    --objective lattice_field \
+    --hidden-dim 512 \
+    --action-horizon 4 \
+    --epochs 150 \
+    --batch-groups 16 \
+    --records-per-group 8 \
+    --lr 0.0003 \
+    --weight-decay 0.01 \
+    --device auto \
+    --seed $SEED \
+    --observation-mode state \
+    --loss-weight bc=1.0 \
+    --loss-weight field_effect=1.5 \
+    --loss-weight field_potential=1.0 \
+    --loss-weight field_preference=0.8 \
+    --loss-weight field_anchor=0.2
+
+echo ""
+echo "✅ Phase A1b enhanced training complete (seed $SEED)"
+echo ""
+echo "Next: Evaluate and compare with 38.43% baseline"

scripts/slurm/phase_a2_train_large_model.sbatch

@@ -0,0 +1,59 @@
+#!/bin/bash
+#SBATCH --job-name=dovla_large_train
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=8
+#SBATCH --gres=gpu:1
+#SBATCH --mem=64000M
+#SBATCH --time=72:00:00
+#SBATCH --output=logs/phase_a2_large_train_%j.out
+#SBATCH --error=logs/phase_a2_large_train_%j.err
+#SBATCH --array=0-2
+
+set -euo pipefail
+
+# Phase A2: Train large capacity model on 10K dataset
+# Target: 40%+ policy success (vs current 29.67%)
+
+PROJECT_DIR="${PROJECT_DIR:-$PWD}"
+cd "$PROJECT_DIR"
+
+source .venv/bin/activate
+
+DATASET="/scratch/$USER/dovla/experiments/maniskill_presuccess_six_task_collection"
+OUT_DIR="/scratch/$USER/dovla/experiments/phase_a2_large_model"
+SEED=$SLURM_ARRAY_TASK_ID
+
+mkdir -p "$OUT_DIR/seed_$SEED" logs
+
+echo "=== Phase A2: Training Large Capacity Model ==="
+echo "Seed: $SEED"
+echo "Hidden dim: 512 (vs current 256)"
+echo "Dataset: 10K groups"
+echo "Target: 40%+ policy success"
+echo ""
+
+python scripts/train_dovla.py \
+    --dataset "$DATASET" \
+    --out "$OUT_DIR/seed_$SEED" \
+    --objective lattice_field \
+    --hidden-dim 512 \
+    --action-horizon 4 \
+    --epochs 100 \
+    --batch-groups 16 \
+    --records-per-group 8 \
+    --lr 0.0003 \
+    --weight-decay 0.01 \
+    --device auto \
+    --seed $SEED \
+    --observation-mode state \
+    --loss-weight bc=1.0 \
+    --loss-weight field_effect=1.0 \
+    --loss-weight field_potential=1.0 \
+    --loss-weight field_preference=0.5 \
+    --loss-weight field_anchor=0.1
+
+echo ""
+echo "✅ Phase A2 complete: Large model trained (seed $SEED)"
+echo ""
+echo "Next: Run phase_a3_eval_large_model.sbatch"

scripts/slurm/phase_a3_eval_large_model.sbatch

@@ -0,0 +1,50 @@
+#!/bin/bash
+#SBATCH --job-name=dovla_large_eval
+#SBATCH --partition=${DOVLA_PARTITION:-compute}
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=8
+#SBATCH --gres=gpu:1
+#SBATCH --mem=32G
+#SBATCH --time=12:00:00
+#SBATCH --output=logs/phase_a3_large_eval_%j.out
+#SBATCH --error=logs/phase_a3_large_eval_%j.err
+#SBATCH --array=0-2
+
+set -euo pipefail
+
+# Phase A3: Evaluate large model with lattice eval + policy rollout
+# Target: Measure improvement over baseline 29.67%
+
+PROJECT_DIR="${PROJECT_DIR:-$PWD}"
+cd "$PROJECT_DIR"
+
+source .venv/bin/activate
+
+CHECKPOINT_DIR="/scratch/$USER/dovla/experiments/phase_a2_large_model/seed_$SLURM_ARRAY_TASK_ID"
+DATASET="/scratch/$USER/dovla/experiments/maniskill_presuccess_six_task_collection"
+OUT_DIR="$CHECKPOINT_DIR"
+
+echo "=== Phase A3: Evaluating Large Model (seed $SLURM_ARRAY_TASK_ID) ==="
+echo ""
+
+# Lattice evaluation
+echo "Running lattice evaluation..."
+python scripts/eval_lattice_checkpoint.py \
+    --checkpoint "$CHECKPOINT_DIR/best.pt" \
+    --dataset "$DATASET" \
+    --out "$OUT_DIR/lattice_eval.json" \
+    --mode field_only \
+    --all-groups
+
+echo ""
+echo "Running policy rollout..."
+python scripts/eval_maniskill_policy_rollout.py \
+    --checkpoint "$CHECKPOINT_DIR/best.pt" \
+    --dataset "$DATASET" \
+    --out "$OUT_DIR/policy_rollout.json" \
+    --num-groups 700 \
+    --mode validation
+
+echo ""
+echo "✅ Phase A3 complete: Evaluation done (seed $SLURM_ARRAY_TASK_ID)"

scripts/slurm/phase_a4_hparam_sweep.sbatch

@@ -0,0 +1,65 @@
+#!/bin/bash
+#SBATCH --job-name=dovla_hparam_sweep
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=8
+#SBATCH --gres=gpu:1
+#SBATCH --mem=64000M
+#SBATCH --time=48:00:00
+#SBATCH --output=logs/phase_a4_hparam_%A_%a.out
+#SBATCH --error=logs/phase_a4_hparam_%A_%a.err
+#SBATCH --array=0-8
+
+set -euo pipefail
+
+# Phase A4: Hyperparameter sweep
+# Grid: 3 LR x 3 hidden_dim = 9 configs
+
+PROJECT_DIR="${PROJECT_DIR:-$PWD}"
+cd "$PROJECT_DIR"
+
+source .venv/bin/activate
+
+DATASET="/scratch/$USER/dovla/experiments/maniskill_presuccess_six_task_collection"
+OUT_ROOT="/scratch/$USER/dovla/experiments/phase_a4_hparam_sweep"
+
+# Hyperparameter grid
+LRS=(0.0001 0.0003 0.001)
+HIDDEN_DIMS=(256 512 1024)
+
+# Map array index to config
+IDX=$SLURM_ARRAY_TASK_ID
+LR_IDX=$((IDX / 3))
+HD_IDX=$((IDX % 3))
+
+LR="${LRS[$LR_IDX]}"
+HIDDEN_DIM="${HIDDEN_DIMS[$HD_IDX]}"
+
+OUT_DIR="$OUT_ROOT/lr${LR}_h${HIDDEN_DIM}"
+mkdir -p "$OUT_DIR" logs
+
+echo "=== Phase A4: Hyperparameter Sweep ==="
+echo "Config $IDX: LR=$LR, Hidden=$HIDDEN_DIM"
+echo ""
+
+python scripts/train_dovla.py \
+    --dataset "$DATASET" \
+    --out "$OUT_DIR" \
+    --objective lattice_field \
+    --hidden-dim "$HIDDEN_DIM" \
+    --epochs 50 \
+    --batch-groups 16 \
+    --lr "$LR" \
+    --device auto \
+    --seed 0
+
+echo ""
+# Quick eval
+python scripts/eval_lattice_checkpoint.py \
+    --checkpoint "$OUT_DIR/best.pt" \
+    --dataset /scratch/$USER/dovla/experiments/maniskill_presuccess_six_task_collection \
+    --out "$OUT_DIR/lattice_eval.json" \
+    --mode field_only \
+    --all-groups
+
+echo "✅ Phase A4 config $IDX complete"

scripts/slurm/phase_a5_horizon_sweep.sbatch

@@ -0,0 +1,63 @@
+#!/bin/bash
+#SBATCH --job-name=dovla_horizon_sweep
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=8
+#SBATCH --gres=gpu:1
+#SBATCH --mem=48000M
+#SBATCH --time=24:00:00
+#SBATCH --output=logs/phase_a5_horizon_%A_%a.out
+#SBATCH --error=logs/phase_a5_horizon_%A_%a.err
+#SBATCH --array=0-3
+
+set -euo pipefail
+
+# Phase A5: Action horizon sweep
+# Test H=4,8,12,16 to see if longer horizons help
+
+PROJECT_DIR="${PROJECT_DIR:-$PWD}"
+cd "$PROJECT_DIR"
+
+source .venv/bin/activate
+
+DATASET="/scratch/$USER/dovla/experiments/maniskill_presuccess_six_task_collection"
+OUT_ROOT="/scratch/$USER/dovla/experiments/phase_a5_horizon_sweep"
+
+HORIZONS=(4 8 12 16)
+HORIZON="${HORIZONS[$SLURM_ARRAY_TASK_ID]}"
+
+OUT_DIR="$OUT_ROOT/h${HORIZON}"
+mkdir -p "$OUT_DIR" logs
+
+echo "=== Phase A5: Action Horizon Sweep ==="
+echo "Horizon: $HORIZON (current baseline: 4)"
+echo ""
+
+python scripts/train_dovla.py \
+    --dataset "$DATASET" \
+    --out "$OUT_DIR" \
+    --objective lattice_field \
+    --hidden-dim 512 \
+    --action-horizon "$HORIZON" \
+    --epochs 50 \
+    --batch-groups 16 \
+    --lr 0.0003 \
+    --device auto \
+    --seed 0
+
+echo ""
+python scripts/eval_lattice_checkpoint.py \
+    --checkpoint "$OUT_DIR/best.pt" \
+    --dataset /scratch/$USER/dovla/experiments/maniskill_presuccess_six_task_collection \
+    --out "$OUT_DIR/lattice_eval.json" \
+    --mode field_only \
+    --all-groups
+
+python scripts/eval_maniskill_policy_rollout.py \
+    --checkpoint "$OUT_DIR/best.pt" \
+    --dataset /scratch/$USER/dovla/experiments/maniskill_presuccess_six_task_collection \
+    --out "$OUT_DIR/policy_rollout.json" \
+    --num-groups 700 \
+    --mode validation
+
+echo "✅ Phase A5 horizon=$HORIZON complete"

scripts/slurm/phase_b_generate_12tasks.sbatch

@@ -0,0 +1,109 @@
+#!/bin/bash
+#SBATCH --job-name=dovla_12task_gen
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=16
+#SBATCH --gres=gpu:1
+#SBATCH --mem=64000M
+#SBATCH --time=72:00:00
+#SBATCH --output=logs/phase_b_12task_gen_%j.out
+#SBATCH --error=logs/phase_b_12task_gen_%j.err
+
+set -euo pipefail
+
+# Phase B Option 1: Generate 12-task ManiSkill collection
+# Fastest option - uses existing infrastructure
+
+PROJECT_DIR="${PROJECT_DIR:-$PWD}"
+cd "$PROJECT_DIR"
+
+source .venv/bin/activate
+
+OUT_DIR="/scratch/$USER/dovla/experiments/phase_b_12task_collection"
+K=16
+STATE_BATCH_SIZE=16
+
+# 12 tasks: 6 existing + 6 new
+declare -A TASK_GROUPS=(
+    # Original 6
+    ["PickCube-v1"]=800
+    ["PushCube-v1"]=800
+    ["PullCube-v1"]=600
+    ["StackCube-v1"]=600
+    ["LiftPegUpright-v1"]=600
+    ["PegInsertionSide-v1"]=600
+
+    # New 6 (TODO: ensure demos exist)
+    ["TurnFaucet-v1"]=500
+    ["OpenDrawer-v1"]=500
+    ["CloseDrawer-v1"]=500
+    ["PlugCharger-v1"]=400
+    ["HangMug-v1"]=400
+    ["PourWater-v1"]=400
+)
+
+mkdir -p "$OUT_DIR" logs
+
+echo "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "="
+echo "Phase B Option 1: 12-Task ManiSkill Collection"
+echo "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "="
+echo ""
+echo "Target: 6,200 groups, 99,200 records (K=$K)"
+echo "Strategy: Expand existing ManiSkill tasks"
+echo ""
+
+# Check if this is just a planning run
+if [ "${DRY_RUN:-0}" = "1" ]; then
+    echo "DRY RUN: Would generate 12 tasks"
+    for TASK in "${!TASK_GROUPS[@]}"; do
+        echo "  $TASK: ${TASK_GROUPS[$TASK]} groups"
+    done
+    exit 0
+fi
+
+# Generate each task
+for TASK in "${!TASK_GROUPS[@]}"; do
+    NUM_GROUPS="${TASK_GROUPS[$TASK]}"
+
+    # Check if already generated
+    if [ -d "$OUT_DIR/${TASK}_k${K}_n${NUM_GROUPS}/merged" ]; then
+        echo "✓ $TASK already exists, skipping"
+        continue
+    fi
+
+    echo "Generating $TASK: $NUM_GROUPS groups..."
+
+    # Use existing generation script
+    python scripts/generate_maniskill_lattice.py \
+        --env-id "$TASK" \
+        --control-mode pd_ee_delta_pose \
+        --out "$OUT_DIR/${TASK}_k${K}_n${NUM_GROUPS}" \
+        --num-groups "$NUM_GROUPS" \
+        --k "$K" \
+        --state-batch-size "$STATE_BATCH_SIZE" \
+        --seed 42 \
+        --pre-success-only \
+        --use-official-demos || {
+            echo "⚠️  $TASK failed (demo might not exist)"
+            continue
+        }
+
+    echo "✅ $TASK complete"
+    echo ""
+done
+
+echo "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "="
+echo "Merging into unified 12-task collection"
+echo "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "="
+
+python scripts/make_cil_collection.py \
+    --source-dirs "$OUT_DIR"/*/merged \
+    --out "$OUT_DIR/merged_12tasks" \
+    --name "phase_b_12task_collection"
+
+echo ""
+echo "✅ Phase B Option 1 complete: 12-task collection ready"
+echo "   Location: $OUT_DIR/merged_12tasks"
+echo ""
+echo "Next: Train on 12 tasks"
+echo "  sbatch scripts/slurm/phase_b_train_12tasks.sbatch"

scripts/slurm/phase_b_train_12tasks.sbatch

@@ -0,0 +1,64 @@
+#!/bin/bash
+#SBATCH --job-name=dovla_12task_train
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=8
+#SBATCH --gres=gpu:1
+#SBATCH --mem=64000M
+#SBATCH --time=96:00:00
+#SBATCH --output=logs/phase_b_12task_train_%A_%a.out
+#SBATCH --error=logs/phase_b_12task_train_%A_%a.err
+#SBATCH --array=0-2
+
+set -euo pipefail
+
+# Phase B: Train on 12-task collection (3 seeds)
+
+PROJECT_DIR="${PROJECT_DIR:-$PWD}"
+cd "$PROJECT_DIR"
+
+source .venv/bin/activate
+
+DATASET="/scratch/$USER/dovla/experiments/phase_b_12task_collection/merged_12tasks"
+OUT_DIR="/scratch/$USER/dovla/experiments/phase_b_12task_model"
+SEED=$SLURM_ARRAY_TASK_ID
+
+mkdir -p "$OUT_DIR/seed_$SEED" logs
+
+echo "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "="
+echo "Phase B: Training on 12-Task Collection"
+echo "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "="
+echo ""
+echo "Seed: $SEED"
+echo "Tasks: 12 (6 existing + 6 new)"
+echo "Groups: ~6,200"
+echo "Hidden dim: 1024 (larger for 12 tasks)"
+echo ""
+
+python scripts/train_dovla.py \
+    --dataset "$DATASET" \
+    --out "$OUT_DIR/seed_$SEED" \
+    --objective lattice_field \
+    --hidden-dim 1024 \
+    --action-horizon 4 \
+    --epochs 100 \
+    --batch-groups 16 \
+    --records-per-group 8 \
+    --lr 0.0003 \
+    --weight-decay 0.01 \
+    --dropout 0.1 \
+    --warmup-steps 1000 \
+    --device auto \
+    --seed $SEED \
+    --observation-mode state \
+    --loss-weight bc=1.0 \
+    --loss-weight field_effect=1.0 \
+    --loss-weight field_utility_regression=1.0 \
+    --loss-weight field_utility_margin=0.5 \
+    --loss-weight field_preference=0.5 \
+    --loss-weight effect_anchor=0.1
+
+echo ""
+echo "✅ Phase B training complete (seed $SEED)"
+echo ""
+echo "Next: Evaluate on held-out tasks"

scripts/slurm/plan_c_generate_10k.sbatch

@@ -0,0 +1,121 @@
+#!/bin/bash
+#SBATCH --job-name=dovla_10k_planc
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=16
+#SBATCH --gres=gpu:1
+#SBATCH --mem=64000M
+#SBATCH --time=96:00:00
+#SBATCH --output=logs/plan_c_10k_gen_%j.out
+#SBATCH --error=logs/plan_c_10k_gen_%j.err
+
+set -euo pipefail
+
+# Plan C: Phase 1B - Generate 10K Groups with Enhanced Sampling
+
+PROJECT_DIR="${PROJECT_DIR:-$PWD}"
+cd "$PROJECT_DIR"
+
+source .venv/bin/activate
+
+OUT_DIR="/scratch/$USER/dovla/experiments/plan_c_10k_enhanced"
+K=16
+STATE_BATCH_SIZE=16
+DEMO_BASE="/scratch/$USER/dovla/maniskill_multitask_demos"
+
+echo "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "="
+echo "Plan C: 10K Generation with Enhanced Sampling"
+echo "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "="
+echo ""
+echo "Target: 48-50%+ policy success"
+echo "Strategy: Maximum quality improvements"
+echo ""
+
+# Task distribution (balanced across difficulty)
+declare -A TASK_GROUPS=(
+    ["PickCube-v1"]=1800
+    ["PushCube-v1"]=1800
+    ["PullCube-v1"]=1600
+    ["StackCube-v1"]=1600
+    ["LiftPegUpright-v1"]=1600
+    ["PegInsertionSide-v1"]=1600
+)
+
+declare -A TASK_DEMOS=(
+    ["PickCube-v1"]="$DEMO_BASE/PickCube-v1/motionplanning/trajectory.h5"
+    ["PushCube-v1"]="$DEMO_BASE/PushCube-v1/motionplanning/trajectory.h5"
+    ["PullCube-v1"]="$DEMO_BASE/PullCube-v1/rl/trajectory.none.pd_ee_delta_pose.physx_cuda.h5"
+    ["StackCube-v1"]="$DEMO_BASE/StackCube-v1/motionplanning/trajectory.h5"
+    ["LiftPegUpright-v1"]="$DEMO_BASE/LiftPegUpright-v1/rl/trajectory.none.pd_ee_delta_pose.physx_cuda.h5"
+    ["PegInsertionSide-v1"]="$DEMO_BASE/PegInsertionSide-v1/motionplanning/trajectory.h5"
+)
+
+TOTAL_GROUPS=0
+for count in "${TASK_GROUPS[@]}"; do
+    TOTAL_GROUPS=$((TOTAL_GROUPS + count))
+done
+
+echo "Task distribution (total: $TOTAL_GROUPS groups):"
+for TASK in "${!TASK_GROUPS[@]}"; do
+    echo "  ${TASK}: ${TASK_GROUPS[$TASK]} groups"
+done
+echo ""
+
+# Generate each task
+for TASK in "${!TASK_GROUPS[@]}"; do
+    NUM_GROUPS="${TASK_GROUPS[$TASK]}"
+    DEMO_PATH="${TASK_DEMOS[$TASK]}"
+    TASK_OUT="$OUT_DIR/${TASK}_k${K}_n${NUM_GROUPS}"
+
+    if [ -d "$TASK_OUT/merged" ]; then
+        echo "✓ $TASK already generated, skipping"
+        continue
+    fi
+
+    if [ ! -f "$DEMO_PATH" ]; then
+        echo "❌ Demo not found: $DEMO_PATH"
+        echo "   Trying alternate location..."
+        # Try RL demos as fallback
+        DEMO_PATH="$DEMO_BASE/${TASK}/rl/trajectory.h5"
+        if [ ! -f "$DEMO_PATH" ]; then
+            echo "   ❌ No demo found, skipping $TASK"
+            continue
+        fi
+    fi
+
+    echo "Generating $TASK: $NUM_GROUPS groups..."
+    echo "  Demo: $DEMO_PATH"
+    echo "  Start: $(date)"
+
+    python scripts/generate_maniskill_lattice.py \
+        --demo "$DEMO_PATH" \
+        --env-id "$TASK" \
+        --control-mode pd_ee_delta_pose \
+        --out "$TASK_OUT" \
+        --num-groups "$NUM_GROUPS" \
+        --k "$K" \
+        --state-batch-size "$STATE_BATCH_SIZE" \
+        --seed 42 \
+        --candidate-mode structured
+
+    echo "  ✅ Complete: $(date)"
+    echo ""
+done
+
+echo "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "="
+echo "Merging all tasks into unified collection"
+echo "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "="
+
+python scripts/make_cil_collection.py \
+    --source-dirs "$OUT_DIR"/*/merged \
+    --out "$OUT_DIR/merged_10k" \
+    --name "plan_c_10k_enhanced"
+
+echo ""
+echo "✅ Plan C Phase 1B Complete!"
+echo ""
+echo "Output: $OUT_DIR/merged_10k"
+echo "Total groups: $TOTAL_GROUPS"
+echo "Total records: $((TOTAL_GROUPS * K))"
+echo ""
+echo "Next: Phase 2A - Attention architecture"

scripts/slurm/prepare_maniskill_baselines.sbatch

@@ -0,0 +1,33 @@
+#!/bin/bash
+#SBATCH --job-name=dovla_ms_baseprep
+#SBATCH --account=def-yalda_cpu
+#SBATCH --partition=cpubase_bycore_b1
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=4
+#SBATCH --mem=24G
+#SBATCH --time=01:00:00
+#SBATCH --output=outputs/hpc/logs/%x_%j.out
+#SBATCH --error=outputs/hpc/logs/%x_%j.err
+
+set -euo pipefail
+
+PROJECT_DIR="${PROJECT_DIR:-$SLURM_SUBMIT_DIR}"
+DATASET="${DATASET:?Set DATASET to the measured CIL collection}"
+OUT_ROOT="${OUT_ROOT:?Set OUT_ROOT for transformed datasets}"
+PYTHON="${PYTHON:-$PROJECT_DIR/.venv/bin/python}"
+
+cd "$PROJECT_DIR"
+mkdir -p "$OUT_ROOT"
+
+"$PYTHON" scripts/prepare_baseline_dataset.py \
+  --dataset "$DATASET" \
+  --baseline expert_only_bc \
+  --out "$OUT_ROOT/expert_only_bc" \
+  --shard-size 2048
+
+"$PYTHON" scripts/prepare_baseline_dataset.py \
+  --dataset "$DATASET" \
+  --baseline label_only_counterfactual \
+  --out "$OUT_ROOT/label_only_counterfactual" \
+  --shard-size 2048

scripts/slurm/render_maniskill_multitask.sbatch

@@ -0,0 +1,33 @@
+#!/bin/bash
+#SBATCH --job-name=dovla_ms_multi_rgb
+#SBATCH --account=def-yalda_cpu
+#SBATCH --partition=cpubase_bycore_b1
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=8
+#SBATCH --mem=24G
+#SBATCH --time=03:00:00
+#SBATCH --array=0-4%2
+#SBATCH --output=outputs/hpc/logs/%x_%A_%a.out
+#SBATCH --error=outputs/hpc/logs/%x_%A_%a.err
+
+set -euo pipefail
+
+PROJECT_DIR="${PROJECT_DIR:-$SLURM_SUBMIT_DIR}"
+MULTITASK_OUT_ROOT="${MULTITASK_OUT_ROOT:?Set MULTITASK_OUT_ROOT}"
+
+case "${SLURM_ARRAY_TASK_ID:-0}" in
+  0) ENV_ID="PushCube-v1" ;;
+  1) ENV_ID="PullCube-v1" ;;
+  2) ENV_ID="StackCube-v1" ;;
+  3) ENV_ID="LiftPegUpright-v1" ;;
+  4) ENV_ID="PegInsertionSide-v1" ;;
+  *) echo "unsupported array index" >&2; exit 2 ;;
+esac
+
+export PROJECT_DIR
+export DATASET="$MULTITASK_OUT_ROOT/$ENV_ID"
+export IMAGE_QUALITY="${IMAGE_QUALITY:-85}"
+export SEED="${SEED:-0}"
+
+exec bash "$PROJECT_DIR/scripts/slurm/render_maniskill_observations.sbatch"

scripts/slurm/render_maniskill_observations.sbatch

@@ -0,0 +1,49 @@
+#!/bin/bash
+#SBATCH --job-name=dovla_ms_render
+#SBATCH --account=def-yalda_cpu
+#SBATCH --partition=cpubase_bycore_b1
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=8
+#SBATCH --mem=24G
+#SBATCH --time=02:00:00
+#SBATCH --output=outputs/hpc/logs/%x_%j.out
+#SBATCH --error=outputs/hpc/logs/%x_%j.err
+
+set -euo pipefail
+
+PROJECT_DIR="${PROJECT_DIR:-$SLURM_SUBMIT_DIR}"
+SCRATCH_ROOT="/scratch/$USER/dovla"
+SIF="$SCRATCH_ROOT/containers/pytorch_2.7.1_cuda12.8.sif"
+PYTHON="$SCRATCH_ROOT/envs/maniskill/bin/python"
+NATIVE_LIBS="$SCRATCH_ROOT/native_libs/lib"
+CPU_RENDER_LIBS="$SCRATCH_ROOT/cpu_render_libs"
+VULKAN_ICD="$CPU_RENDER_LIBS/share/vulkan/icd.d/lvp_icd.x86_64.json"
+
+DATASET="${DATASET:?Set DATASET to a generated ManiSkill CIL directory}"
+IMAGE_QUALITY="${IMAGE_QUALITY:-90}"
+SEED="${SEED:-0}"
+OVERWRITE="${OVERWRITE:-0}"
+RUNTIME_DIR="/tmp/$USER/dovla-render-$SLURM_JOB_ID"
+CACHE_DIR="/tmp/$USER/dovla-render-cache-$SLURM_JOB_ID"
+
+module load StdEnv/2023 apptainer/1.4.5
+cd "$PROJECT_DIR"
+mkdir -p "$RUNTIME_DIR" "$CACHE_DIR"
+chmod 700 "$RUNTIME_DIR"
+
+ARGS=(
+  --dataset "$DATASET"
+  --render-backend cpu
+  --image-quality "$IMAGE_QUALITY"
+  --seed "$SEED"
+)
+if [[ "$OVERWRITE" == "1" ]]; then
+  ARGS+=(--overwrite)
+fi
+
+apptainer exec \
+  --env "LD_LIBRARY_PATH=$CPU_RENDER_LIBS/lib:$NATIVE_LIBS,XDG_RUNTIME_DIR=$RUNTIME_DIR,MESA_SHADER_CACHE_DIR=$CACHE_DIR,LIBGL_ALWAYS_SOFTWARE=1,LP_NUM_THREADS=4,VK_ICD_FILENAMES=$VULKAN_ICD,VK_DRIVER_FILES=$VULKAN_ICD,OMP_NUM_THREADS=1,OPENBLAS_NUM_THREADS=1,MKL_NUM_THREADS=1" \
+  "$SIF" "$PYTHON" scripts/render_maniskill_observations.py "${ARGS[@]}"
+
+rm -rf "$RUNTIME_DIR" "$CACHE_DIR"

scripts/slurm/run_external_vla_baseline.sbatch

@@ -0,0 +1,51 @@
+#!/bin/bash
+#SBATCH --job-name=dovla_ext_vla
+#SBATCH --account=def-yalda_gpu
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=8
+#SBATCH --gres=gpu:nvidia_h100_80gb_hbm3_1g.10gb:1
+#SBATCH --mem=40G
+#SBATCH --time=04:00:00
+#SBATCH --output=outputs/hpc/logs/%x_%j.out
+#SBATCH --error=outputs/hpc/logs/%x_%j.err
+
+set -euo pipefail
+
+PROJECT_DIR="${PROJECT_DIR:-$SLURM_SUBMIT_DIR}"
+MODEL_FAMILY="${MODEL_FAMILY:-smolvla}"
+DATASET="${DATASET:?Set DATASET to the held-out DoVLA-CIL dataset}"
+OUT="${OUT:?Set OUT to a run directory}"
+CHECKPOINT="${CHECKPOINT:-}"
+ADAPTER_ENTRYPOINT="${ADAPTER_ENTRYPOINT:-}"
+ADAPTER_CONFIG="${ADAPTER_CONFIG:-}"
+PYTHON="${PYTHON:-python}"
+DRY_RUN="${DRY_RUN:-0}"
+
+cd "$PROJECT_DIR"
+mkdir -p "$OUT" outputs/hpc/logs
+
+ARGS=(
+  scripts/run_external_vla_baseline.py
+  --model-family "$MODEL_FAMILY"
+  --dataset "$DATASET"
+  --out "$OUT"
+  --python "$PYTHON"
+)
+
+if [[ -n "$CHECKPOINT" ]]; then
+  ARGS+=(--checkpoint "$CHECKPOINT")
+fi
+if [[ -n "$ADAPTER_ENTRYPOINT" ]]; then
+  ARGS+=(--adapter-entrypoint "$ADAPTER_ENTRYPOINT")
+fi
+if [[ -n "$ADAPTER_CONFIG" ]]; then
+  ARGS+=(--adapter-config "$ADAPTER_CONFIG")
+fi
+if [[ "$DRY_RUN" == "1" ]]; then
+  ARGS+=(--dry-run)
+else
+  ARGS+=(--require-ready)
+fi
+
+"$PYTHON" "${ARGS[@]}"

scripts/slurm/run_scaling.sbatch

@@ -0,0 +1,42 @@
+#!/bin/bash
+#SBATCH --job-name=${DOVLA_JOB_NAME:-dovla_scaling}
+#SBATCH --partition=${DOVLA_PARTITION:-gpu}
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=${DOVLA_CPUS_PER_TASK:-8}
+#SBATCH --gres=gpu:${DOVLA_GPUS_PER_TASK:-1}
+#SBATCH --mem=${DOVLA_MEM:-64G}
+#SBATCH --time=${DOVLA_TIME:-24:00:00}
+#SBATCH --output=${DOVLA_LOG_DIR:-logs/slurm}/%x_%j.out
+#SBATCH --error=${DOVLA_LOG_DIR:-logs/slurm}/%x_%j.err
+
+set -euo pipefail
+
+PROJECT_DIR="${PROJECT_DIR:-$PWD}"
+VENV_PATH="${VENV_PATH:-$PROJECT_DIR/.venv}"
+BACKEND="${BACKEND:-toy}"
+TASKS="${TASKS:-builtins}"
+OUT_DIR="${OUT_DIR:-$PROJECT_DIR/runs/scaling_toy}"
+TOTAL_RECORDS="${TOTAL_RECORDS:-4096}"
+K_VALUES="${K_VALUES:-1,2,4,8,16,32}"
+EPOCHS="${EPOCHS:-3}"
+SEED="${SEED:-0}"
+DEVICE="${DEVICE:-auto}"
+
+mkdir -p "${DOVLA_LOG_DIR:-logs/slurm}" "$OUT_DIR"
+cd "$PROJECT_DIR"
+
+if [ -f "$VENV_PATH/bin/activate" ]; then
+  # shellcheck disable=SC1091
+  source "$VENV_PATH/bin/activate"
+fi
+
+python scripts/run_scaling.py \
+  --backend "$BACKEND" \
+  --tasks "$TASKS" \
+  --out "$OUT_DIR" \
+  --total-records "$TOTAL_RECORDS" \
+  --k-values "$K_VALUES" \
+  --epochs "$EPOCHS" \
+  --seed "$SEED" \
+  --device "$DEVICE"

scripts/slurm/run_smolvla_cil_baseline.sbatch

@@ -0,0 +1,48 @@
+#!/bin/bash
+#SBATCH --job-name=dovla_smolvla_cil
+#SBATCH --account=def-yalda_gpu
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=8
+#SBATCH --gres=gpu:nvidia_h100_80gb_hbm3_3g.40gb:1
+#SBATCH --mem=40G
+#SBATCH --time=02:00:00
+#SBATCH --output=outputs/hpc/logs/%x_%j.out
+#SBATCH --error=outputs/hpc/logs/%x_%j.err
+
+set -euo pipefail
+
+PROJECT_DIR="${PROJECT_DIR:-$SLURM_SUBMIT_DIR}"
+SCRATCH_ROOT="${SCRATCH_ROOT:-/scratch/$USER/dovla}"
+CONTAINER="${CONTAINER:-$SCRATCH_ROOT/containers/pytorch_2.7.1_cuda12.8.sif}"
+PYTHON="${PYTHON:-$SCRATCH_ROOT/envs/smolvla/bin/python}"
+CHECKPOINT="${CHECKPOINT:-$SCRATCH_ROOT/models/smolvla_base-c83c316}"
+DATASET="${DATASET:-$SCRATCH_ROOT/experiments/maniskill_presuccess_six_task_collection}"
+ADAPTER_CONFIG="${ADAPTER_CONFIG:-$PROJECT_DIR/configs/external/smolvla_cil_smoke.json}"
+OUT="${OUT:-$SCRATCH_ROOT/experiments/smolvla_cil_smoke}"
+CONTAINER_ADAPTER_CONFIG="$ADAPTER_CONFIG"
+if [[ "$ADAPTER_CONFIG" == "$PROJECT_DIR/"* ]]; then
+  CONTAINER_ADAPTER_CONFIG="/workspace/${ADAPTER_CONFIG#"$PROJECT_DIR/"}"
+fi
+
+cd "$PROJECT_DIR"
+mkdir -p "$OUT" outputs/hpc/logs
+module load StdEnv/2023 apptainer/1.4.5
+
+apptainer exec \
+  --nv \
+  -B "$SCRATCH_ROOT:$SCRATCH_ROOT" \
+  -B "$PROJECT_DIR:/workspace" \
+  --env \
+  "PYTHONNOUSERSITE=1,HF_HUB_OFFLINE=1,TRANSFORMERS_OFFLINE=1,SCRATCH_ROOT=$SCRATCH_ROOT" \
+  "$CONTAINER" \
+  "$PYTHON" /workspace/scripts/run_external_vla_baseline.py \
+    --model-family smolvla \
+    --checkpoint "$CHECKPOINT" \
+    --dataset "$DATASET" \
+    --out "$OUT" \
+    --python "$PYTHON" \
+    --adapter-entrypoint \
+    dovla_cil.eval.smolvla_cil_baseline:run_smolvla_cil_baseline \
+    --adapter-config "$CONTAINER_ADAPTER_CONFIG" \
+    --require-ready

scripts/slurm/smoke_smolvla_checkpoint.sbatch

@@ -0,0 +1,42 @@
+#!/bin/bash
+#SBATCH --job-name=dovla_smolvla_smoke
+#SBATCH --account=def-yalda_gpu
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=4
+#SBATCH --gres=gpu:nvidia_h100_80gb_hbm3_1g.10gb:1
+#SBATCH --mem=24G
+#SBATCH --time=00:30:00
+#SBATCH --output=outputs/hpc/logs/%x_%j.out
+#SBATCH --error=outputs/hpc/logs/%x_%j.err
+
+set -euo pipefail
+
+PROJECT_DIR="${PROJECT_DIR:-$SLURM_SUBMIT_DIR}"
+SCRATCH_ROOT="${SCRATCH_ROOT:-/scratch/$USER/dovla}"
+CONTAINER="${CONTAINER:-$SCRATCH_ROOT/containers/pytorch_2.7.1_cuda12.8.sif}"
+PYTHON="${PYTHON:-$SCRATCH_ROOT/envs/smolvla/bin/python}"
+CHECKPOINT="${CHECKPOINT:-$SCRATCH_ROOT/models/smolvla_base-c83c316}"
+VLM_REVISION="${VLM_REVISION:-7b375e1b73b11138ff12fe22c8f2822d8fe03467}"
+VLM_METADATA="${VLM_METADATA:-$SCRATCH_ROOT/models/SmolVLM2-500M-Video-Instruct-metadata-$VLM_REVISION}"
+OUT="${OUT:-$PROJECT_DIR/outputs/external_vla_smolvla_checkpoint_smoke.json}"
+CONTAINER_OUT="$OUT"
+if [[ "$OUT" == "$PROJECT_DIR/"* ]]; then
+  CONTAINER_OUT="/workspace/${OUT#"$PROJECT_DIR/"}"
+fi
+
+cd "$PROJECT_DIR"
+mkdir -p "$(dirname "$OUT")" outputs/hpc/logs
+module load StdEnv/2023 apptainer/1.4.5
+
+apptainer exec \
+  --nv \
+  -B "$SCRATCH_ROOT:$SCRATCH_ROOT" \
+  -B "$PROJECT_DIR:/workspace" \
+  --env PYTHONNOUSERSITE=1,HF_HUB_OFFLINE=1,TRANSFORMERS_OFFLINE=1 \
+  "$CONTAINER" \
+  "$PYTHON" /workspace/scripts/smoke_smolvla_checkpoint.py \
+    --checkpoint "$CHECKPOINT" \
+    --vlm-metadata "$VLM_METADATA" \
+    --out "$CONTAINER_OUT" \
+    --device cuda

scripts/slurm/train_attention_model.sbatch

@@ -0,0 +1,62 @@
+#!/bin/bash
+#SBATCH --job-name=dovla_attention
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=8
+#SBATCH --gres=gpu:1
+#SBATCH --mem=64000M
+#SBATCH --time=48:00:00
+#SBATCH --output=logs/attention_train_%A_%a.out
+#SBATCH --error=logs/attention_train_%A_%a.err
+#SBATCH --array=0-2
+
+set -euo pipefail
+
+# CVPR-Ready: DoVLA-Attention Architecture
+# Single principled contribution: Transformer attention for action comparison
+
+PROJECT_DIR="${PROJECT_DIR:-$PWD}"
+cd "$PROJECT_DIR"
+
+source .venv/bin/activate
+
+DATASET="/scratch/$USER/dovla/experiments/maniskill_presuccess_six_task_collection"
+OUT_DIR="/scratch/$USER/dovla/experiments/cvpr_attention_model"
+SEED=$SLURM_ARRAY_TASK_ID
+
+mkdir -p "$OUT_DIR/seed_$SEED" logs
+
+echo "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "="
+echo "CVPR Experiment: DoVLA-Attention Architecture"
+echo "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "="
+echo ""
+echo "Method: Transformer-based attention for action comparison"
+echo "Contribution: Cross-attention + Self-attention + Pairwise head"
+echo "Dataset: 3,500 groups (SAME as baseline for fair comparison)"
+echo "Seed: $SEED"
+echo ""
+echo "Expected: 42-44% success (vs 38.43% MLP baseline)"
+echo ""
+
+# Train with attention architecture
+python scripts/train_dovla_attention.py \
+    --dataset "$DATASET" \
+    --out "$OUT_DIR/seed_$SEED" \
+    --model attention \
+    --hidden-dim 256 \
+    --n-heads 4 \
+    --n-layers 2 \
+    --action-horizon 4 \
+    --epochs 50 \
+    --batch-groups 16 \
+    --records-per-group 8 \
+    --lr 0.0003 \
+    --weight-decay 0.01 \
+    --device auto \
+    --seed $SEED \
+    --observation-mode state
+
+echo ""
+echo "✅ DoVLA-Attention training complete (seed $SEED)"
+echo ""
+echo "Next: Evaluate and compare with MLP baseline (38.43%)"

scripts/slurm/train_dovla.sbatch

@@ -0,0 +1,44 @@
+#!/bin/bash
+#SBATCH --job-name=${DOVLA_JOB_NAME:-dovla_train}
+#SBATCH --partition=${DOVLA_PARTITION:-gpu}
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=${DOVLA_CPUS_PER_TASK:-8}
+#SBATCH --gres=gpu:${DOVLA_GPUS_PER_TASK:-1}
+#SBATCH --mem=${DOVLA_MEM:-64G}
+#SBATCH --time=${DOVLA_TIME:-24:00:00}
+#SBATCH --output=${DOVLA_LOG_DIR:-logs/slurm}/%x_%j.out
+#SBATCH --error=${DOVLA_LOG_DIR:-logs/slurm}/%x_%j.err
+
+set -euo pipefail
+
+PROJECT_DIR="${PROJECT_DIR:-$PWD}"
+VENV_PATH="${VENV_PATH:-$PROJECT_DIR/.venv}"
+DATASET="${DATASET:-$PROJECT_DIR/data/cil_toy}"
+OUT_DIR="${OUT_DIR:-$PROJECT_DIR/runs/dovla_toy}"
+EPOCHS="${EPOCHS:-5}"
+BATCH_GROUPS="${BATCH_GROUPS:-8}"
+RECORDS_PER_GROUP="${RECORDS_PER_GROUP:-8}"
+HIDDEN_DIM="${HIDDEN_DIM:-256}"
+LR="${LR:-0.001}"
+DEVICE="${DEVICE:-auto}"
+SEED="${SEED:-0}"
+
+mkdir -p "${DOVLA_LOG_DIR:-logs/slurm}" "$OUT_DIR"
+cd "$PROJECT_DIR"
+
+if [ -f "$VENV_PATH/bin/activate" ]; then
+  # shellcheck disable=SC1091
+  source "$VENV_PATH/bin/activate"
+fi
+
+python scripts/train_dovla.py \
+  --dataset "$DATASET" \
+  --out "$OUT_DIR" \
+  --epochs "$EPOCHS" \
+  --batch-groups "$BATCH_GROUPS" \
+  --records-per-group "$RECORDS_PER_GROUP" \
+  --hidden-dim "$HIDDEN_DIM" \
+  --lr "$LR" \
+  --device "$DEVICE" \
+  --seed "$SEED"

scripts/slurm/train_enhanced_model.sbatch

@@ -0,0 +1,63 @@
+#!/bin/bash
+#SBATCH --job-name=dovla_enhanced
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=8
+#SBATCH --gres=gpu:1
+#SBATCH --mem=64000M
+#SBATCH --time=48:00:00
+#SBATCH --output=logs/enhanced_train_%A_%a.out
+#SBATCH --error=logs/enhanced_train_%A_%a.err
+#SBATCH --array=0-2
+
+set -euo pipefail
+
+# DoVLA-Attention-Enhanced: SOTA Architecture for CVPR
+# Hierarchical attention + Graph NN + Contrastive + Task-adaptive
+
+PROJECT_DIR="${PROJECT_DIR:-$PWD}"
+cd "$PROJECT_DIR"
+
+source .venv/bin/activate
+
+DATASET="/scratch/$USER/dovla/experiments/maniskill_presuccess_six_task_collection"
+OUT_DIR="/scratch/$USER/dovla/experiments/cvpr_enhanced_model"
+SEED=$SLURM_ARRAY_TASK_ID
+
+mkdir -p "$OUT_DIR/seed_$SEED" logs
+
+echo "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "="
+echo "DoVLA-Attention-Enhanced: SOTA Training"
+echo "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "="
+echo ""
+echo "Architecture Components:"
+echo "  1. Hierarchical Attention (local + global)"
+echo "  2. Graph Neural Network (action relationships)"
+echo "  3. Contrastive Learning (better embeddings)"
+echo "  4. Task-Adaptive Layers (multi-task)"
+echo ""
+echo "Dataset: 3,500 groups (fair comparison)"
+echo "Seed: $SEED"
+echo ""
+echo "Expected: 44-47% success (vs 38.43% baseline)"
+echo "Improvement: +5.5-8.5%"
+echo ""
+
+python scripts/train_dovla_enhanced.py \
+    --dataset "$DATASET" \
+    --out "$OUT_DIR/seed_$SEED" \
+    --hidden-dim 256 \
+    --n-heads 4 \
+    --n-layers 3 \
+    --epochs 50 \
+    --batch-size 16 \
+    --lr 0.0003 \
+    --weight-decay 0.01 \
+    --contrastive-weight 0.1 \
+    --seed $SEED \
+    --device auto
+
+echo ""
+echo "✅ Enhanced training complete (seed $SEED)"
+echo ""
+echo "Next: Evaluate and compare with baseline"

scripts/slurm/train_h16_policy.sbatch

@@ -0,0 +1,54 @@
+#!/bin/bash
+#SBATCH --job-name=train_h16_policy
+#SBATCH --account=def-yalda_gpu
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=8
+#SBATCH --gres=gpu:1
+#SBATCH --mem=48G
+#SBATCH --time=04:00:00
+#SBATCH --output=logs/train_h16_%A_%a.out
+#SBATCH --error=logs/train_h16_%A_%a.err
+#SBATCH --array=0-2
+
+set -euo pipefail
+
+# Train policy on h=16 collection (oracle 94.76%)
+# Expected: val top-1 ~85-90%, online rollout 55-70%+
+
+PROJECT_DIR="${PROJECT_DIR:-$SLURM_SUBMIT_DIR}"
+cd "$PROJECT_DIR"
+
+source .venv/bin/activate
+
+SEED=$SLURM_ARRAY_TASK_ID
+DATASET="/scratch/$USER/dovla/experiments/six_task_h16_collection"
+OUT_DIR="/scratch/$USER/dovla/experiments/h16_policy_runs/seed_$SEED"
+
+mkdir -p "$OUT_DIR" logs
+
+echo "=================================================="
+echo "Training Policy on h=16 Collection"
+echo "Seed: $SEED"
+echo "Dataset: $DATASET"
+echo "Expected oracle: 94.76%"
+echo "Expected val top-1: 85-90%"
+echo "=================================================="
+
+python scripts/train_hybrid_direct.py \
+    --dataset "$DATASET" \
+    --out "$OUT_DIR" \
+    --d-model 256 \
+    --n-heads 8 \
+    --n-layers 4 \
+    --d-ff 1024 \
+    --epochs 50 \
+    --batch-size 128 \
+    --lr 3e-4 \
+    --warmup-steps 500 \
+    --seed "$SEED" \
+    --device cuda
+
+echo ""
+echo "✅ Training complete for seed $SEED"
+echo "Best checkpoint: $OUT_DIR/best.pt"

scripts/slurm/train_hybrid_direct.sbatch

@@ -0,0 +1,65 @@
+#!/bin/bash
+#SBATCH --job-name=hybrid_direct
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=8
+#SBATCH --gres=gpu:1
+#SBATCH --mem=64000M
+#SBATCH --time=48:00:00
+#SBATCH --output=logs/hybrid_direct_%A_%a.out
+#SBATCH --error=logs/hybrid_direct_%A_%a.err
+#SBATCH --array=0-2
+
+set -euo pipefail
+
+# DoVLA-Hybrid: DIRECT Scoring (NOT Pairwise)
+# Expected: 45-48% baseline (vs 37% pairwise)
+
+PROJECT_DIR="${PROJECT_DIR:-$PWD}"
+cd "$PROJECT_DIR"
+
+source .venv/bin/activate
+
+DATASET="/scratch/$USER/dovla/experiments/maniskill_presuccess_six_task_collection"
+OUT_DIR="/scratch/$USER/dovla/experiments/cvpr_hybrid_direct_model"
+SEED=$SLURM_ARRAY_TASK_ID
+
+mkdir -p "$OUT_DIR/seed_$SEED" logs
+
+echo "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "="
+echo "DoVLA-Hybrid: DIRECT Scoring (FIXED!)"
+echo "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "="
+echo ""
+echo "KEY IMPROVEMENT:"
+echo "  OLD: Pairwise ranking → aggregate → 37%"
+echo "  NEW: Direct scoring → 45-48%"
+echo ""
+echo "Approach:"
+echo "  - Predict reward(action) directly"
+echo "  - Predict success(action) directly"
+echo "  - Select: argmax(success_prob * reward)"
+echo ""
+echo "Expected: 45-48% WITHOUT language"
+echo "Then +language: 55-60% final"
+echo ""
+echo "Seed: $SEED"
+echo ""
+
+python scripts/train_hybrid_direct.py \
+    --dataset "$DATASET" \
+    --out "$OUT_DIR/seed_$SEED" \
+    --d-model 256 \
+    --n-heads 8 \
+    --n-layers 3 \
+    --d-ff 1024 \
+    --epochs 50 \
+    --batch-size 16 \
+    --lr 0.001 \
+    --weight-decay 0.01 \
+    --warmup-steps 500 \
+    --seed $SEED \
+    --device auto
+
+echo ""
+echo "✅ Hybrid training complete (seed $SEED)"
+echo ""

scripts/slurm/train_maniskill_baseline_array.sbatch

@@ -0,0 +1,84 @@
+#!/bin/bash
+#SBATCH --job-name=dovla_ms_baseline
+#SBATCH --account=def-yalda_gpu
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=4
+#SBATCH --gres=gpu:nvidia_h100_80gb_hbm3_1g.10gb:1
+#SBATCH --mem=28G
+#SBATCH --time=02:00:00
+#SBATCH --array=0-2%3
+#SBATCH --output=outputs/hpc/logs/%x_%A_%a.out
+#SBATCH --error=outputs/hpc/logs/%x_%A_%a.err
+
+set -euo pipefail
+
+PROJECT_DIR="${PROJECT_DIR:-$SLURM_SUBMIT_DIR}"
+BASELINE="${BASELINE:?Set BASELINE}"
+DATASET="${DATASET:?Set DATASET}"
+RUN_ROOT="${RUN_ROOT:?Set RUN_ROOT}"
+PYTHON="${PYTHON:-$PROJECT_DIR/.venv/bin/python}"
+SEED="${SLURM_ARRAY_TASK_ID:-0}"
+EPOCHS="${EPOCHS:-50}"
+RECORDS_PER_GROUP="${RECORDS_PER_GROUP:-16}"
+PAIR_SCOPE="same_state"
+LOSS_ARGS=()
+
+case "$BASELINE" in
+  cross_state_negatives)
+    PAIR_SCOPE="cross_state"
+    ;;
+  random_negatives)
+    ;;
+  world_model_auxiliary|no_rank_regret)
+    LOSS_ARGS+=(--loss-weight rank=0 --loss-weight regret=0)
+    ;;
+  no_effect_head)
+    LOSS_ARGS+=(--loss-weight effect=0)
+    ;;
+  label_only_counterfactual)
+    LOSS_ARGS+=(--loss-weight effect=0)
+    ;;
+  expert_only_bc)
+    RECORDS_PER_GROUP=1
+    LOSS_ARGS+=(
+      --loss-weight effect=0
+      --loss-weight progress=0
+      --loss-weight rank=0
+      --loss-weight regret=0
+    )
+    ;;
+  *)
+    echo "unsupported baseline: $BASELINE" >&2
+    exit 2
+    ;;
+esac
+
+OUT_DIR="$RUN_ROOT/$BASELINE/seed_$SEED"
+cd "$PROJECT_DIR"
+mkdir -p "$OUT_DIR"
+export OMP_NUM_THREADS=1
+export OPENBLAS_NUM_THREADS=1
+export MKL_NUM_THREADS=1
+export DOVLA_TORCH_THREADS=1
+
+"$PYTHON" scripts/train_dovla.py \
+  --dataset "$DATASET" \
+  --out "$OUT_DIR" \
+  --epochs "$EPOCHS" \
+  --batch-groups 32 \
+  --records-per-group "$RECORDS_PER_GROUP" \
+  --pair-count-per-group 32 \
+  --hidden-dim 256 \
+  --obs-dim 96 \
+  --lang-dim 64 \
+  --action-dim 8 \
+  --action-horizon 4 \
+  --effect-dim 32 \
+  --lr 0.001 \
+  --device cuda \
+  --seed "$SEED" \
+  --val-fraction 0.2 \
+  --objective legacy \
+  --pair-scope "$PAIR_SCOPE" \
+  "${LOSS_ARGS[@]}"

scripts/slurm/train_maniskill_collection_array.sbatch

@@ -0,0 +1,114 @@
+#!/bin/bash
+#SBATCH --job-name=dovla_ms_multi_train
+#SBATCH --account=def-yalda_gpu
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=4
+#SBATCH --gres=gpu:nvidia_h100_80gb_hbm3_1g.10gb:1
+#SBATCH --mem=28G
+#SBATCH --time=02:00:00
+#SBATCH --array=0-5%6
+#SBATCH --output=outputs/hpc/logs/%x_%A_%a.out
+#SBATCH --error=outputs/hpc/logs/%x_%A_%a.err
+
+set -euo pipefail
+
+PROJECT_DIR="${PROJECT_DIR:-$SLURM_SUBMIT_DIR}"
+DATASET="${DATASET:?Set DATASET to a CIL collection}"
+RUN_ROOT="${RUN_ROOT:?Set RUN_ROOT}"
+PYTHON="${PYTHON:-$PROJECT_DIR/.venv/bin/python}"
+BACKBONE="${BACKBONE:-native}"
+OBSERVATION_MODE="${OBSERVATION_MODE:-state}"
+BACKBONE_MODEL="${BACKBONE_MODEL:-}"
+BACKBONE_FEATURE_CACHE="${BACKBONE_FEATURE_CACHE:-}"
+BACKBONE_FEATURE_BATCH_SIZE="${BACKBONE_FEATURE_BATCH_SIZE:-64}"
+TASK_INDEX="${SLURM_ARRAY_TASK_ID:-0}"
+OBJECTIVE_MODE="${OBJECTIVE_MODE:-paired}"
+EPOCHS="${EPOCHS:-50}"
+BATCH_GROUPS="${BATCH_GROUPS:-32}"
+HIDDEN_DIM="${HIDDEN_DIM:-256}"
+if [[ "$OBJECTIVE_MODE" == "field_only" ]]; then
+  SEED="$TASK_INDEX"
+  OBJECTIVE="${OBJECTIVE:-lattice_field}"
+elif [[ "$OBJECTIVE_MODE" == "paired" ]]; then
+  SEED="$((TASK_INDEX / 2))"
+  if (( TASK_INDEX % 2 == 0 )); then
+    OBJECTIVE="lattice_field"
+  else
+    OBJECTIVE="legacy"
+  fi
+else
+  echo "OBJECTIVE_MODE must be paired or field_only" >&2
+  exit 2
+fi
+OUT_DIR="$RUN_ROOT/$OBJECTIVE/seed_$SEED"
+
+cd "$PROJECT_DIR"
+mkdir -p "$OUT_DIR"
+export OMP_NUM_THREADS=1
+export OPENBLAS_NUM_THREADS=1
+export MKL_NUM_THREADS=1
+export DOVLA_TORCH_THREADS=1
+
+if [[ "$BACKBONE" == "clip" ]]; then
+  [[ "$OBSERVATION_MODE" == "rgb" ]] || { echo "CLIP requires OBSERVATION_MODE=rgb" >&2; exit 2; }
+  [[ -n "$BACKBONE_MODEL" ]] || { echo "Set BACKBONE_MODEL for CLIP" >&2; exit 2; }
+  [[ -n "$BACKBONE_FEATURE_CACHE" ]] || { echo "Set BACKBONE_FEATURE_CACHE for CLIP" >&2; exit 2; }
+fi
+if [[ "$BACKBONE" == "clip" || "$OBSERVATION_MODE" == "rgb" ]]; then
+  SCRATCH_ROOT="/scratch/$USER/dovla"
+  SIF="${SIF:-$SCRATCH_ROOT/containers/pytorch_2.7.1_cuda12.8.sif}"
+  CONTAINER_PYTHON="${CONTAINER_PYTHON:-$SCRATCH_ROOT/envs/maniskill/bin/python}"
+  module load StdEnv/2023 apptainer/1.4.5
+  PYTHON_COMMAND=(
+    apptainer exec --nv
+    --env "OMP_NUM_THREADS=1,OPENBLAS_NUM_THREADS=1,MKL_NUM_THREADS=1,DOVLA_TORCH_THREADS=1,TRANSFORMERS_OFFLINE=1,HF_HUB_OFFLINE=1"
+    -B "$PROJECT_DIR:$PROJECT_DIR"
+    -B "/scratch/$USER:/scratch/$USER"
+    "$SIF" "$CONTAINER_PYTHON"
+  )
+else
+  PYTHON_COMMAND=("$PYTHON")
+fi
+
+"${PYTHON_COMMAND[@]}" - <<PY
+from dovla_cil.data.datasets import CILDataset
+import torch
+
+dataset = CILDataset("$DATASET")
+assert len(dataset.group_ids) == int("${EXPECTED_GROUPS:-3500}"), len(dataset.group_ids)
+assert len(dataset) == int("${EXPECTED_RECORDS:-56000}"), len(dataset)
+assert torch.cuda.is_available()
+print("objective=$OBJECTIVE seed=$SEED groups=", len(dataset.group_ids), "records=", len(dataset), torch.cuda.get_device_name(0))
+PY
+
+BACKBONE_ARGS=(--backbone "$BACKBONE")
+if [[ "$BACKBONE" == "clip" ]]; then
+  BACKBONE_ARGS+=(
+    --backbone-model "$BACKBONE_MODEL"
+    --backbone-feature-cache "$BACKBONE_FEATURE_CACHE"
+    --backbone-feature-batch-size "$BACKBONE_FEATURE_BATCH_SIZE"
+  )
+fi
+
+"${PYTHON_COMMAND[@]}" scripts/train_dovla.py \
+  --dataset "$DATASET" \
+  --out "$OUT_DIR" \
+  --epochs "$EPOCHS" \
+  --batch-groups "$BATCH_GROUPS" \
+  --records-per-group 16 \
+  --pair-count-per-group 32 \
+  --hidden-dim "$HIDDEN_DIM" \
+  --obs-dim 96 \
+  --observation-mode "$OBSERVATION_MODE" \
+  --lang-dim 64 \
+  --action-dim 8 \
+  --action-horizon 4 \
+  --effect-dim 32 \
+  --lr 0.001 \
+  --device cuda \
+  --seed "$SEED" \
+  --val-fraction 0.2 \
+  --objective "$OBJECTIVE" \
+  --lattice-neighbors 32 \
+  "${BACKBONE_ARGS[@]}"

scripts/slurm/train_maniskill_collection_cpu_array.sbatch

@@ -0,0 +1,61 @@
+#!/bin/bash
+#SBATCH --job-name=dovla_ms_cpu_train
+#SBATCH --account=def-yalda_cpu
+#SBATCH --partition=cpubase_bycore_b2
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=8
+#SBATCH --mem=32G
+#SBATCH --time=03:00:00
+#SBATCH --array=0-2%3
+#SBATCH --output=outputs/hpc/logs/%x_%A_%a.out
+#SBATCH --error=outputs/hpc/logs/%x_%A_%a.err
+
+set -euo pipefail
+
+PROJECT_DIR="${PROJECT_DIR:-$SLURM_SUBMIT_DIR}"
+DATASET="${DATASET:?Set DATASET to a CIL collection}"
+RUN_ROOT="${RUN_ROOT:?Set RUN_ROOT}"
+PYTHON="${PYTHON:-$PROJECT_DIR/.venv/bin/python}"
+OBJECTIVE="${OBJECTIVE:-lattice_field}"
+EPOCHS="${EPOCHS:-50}"
+HIDDEN_DIM="${HIDDEN_DIM:-256}"
+TASK_INDEX="${SLURM_ARRAY_TASK_ID:-0}"
+SEED="${SEED_OVERRIDE:-$TASK_INDEX}"
+OUT_DIR="$RUN_ROOT/$OBJECTIVE/seed_$SEED"
+
+cd "$PROJECT_DIR"
+mkdir -p outputs/hpc/logs "$OUT_DIR"
+export OMP_NUM_THREADS=1
+export OPENBLAS_NUM_THREADS=1
+export MKL_NUM_THREADS=1
+export DOVLA_TORCH_THREADS=1
+
+"$PYTHON" - <<PY
+from dovla_cil.data.datasets import CILDataset
+
+dataset = CILDataset("$DATASET")
+assert len(dataset.group_ids) == int("${EXPECTED_GROUPS:-3500}"), len(dataset.group_ids)
+assert len(dataset) == int("${EXPECTED_RECORDS:-56000}"), len(dataset)
+print("objective=$OBJECTIVE seed=$SEED device=cpu groups=", len(dataset.group_ids), "records=", len(dataset))
+PY
+
+"$PYTHON" scripts/train_dovla.py \
+  --dataset "$DATASET" \
+  --out "$OUT_DIR" \
+  --epochs "$EPOCHS" \
+  --batch-groups 32 \
+  --records-per-group 16 \
+  --pair-count-per-group 32 \
+  --hidden-dim "$HIDDEN_DIM" \
+  --obs-dim 96 \
+  --lang-dim 64 \
+  --action-dim 8 \
+  --action-horizon 4 \
+  --effect-dim 32 \
+  --lr 0.001 \
+  --device cpu \
+  --seed "$SEED" \
+  --val-fraction 0.2 \
+  --objective "$OBJECTIVE" \
+  --lattice-neighbors 32

scripts/slurm/train_maniskill_debug.sbatch

@@ -0,0 +1,86 @@
+#!/bin/bash
+#SBATCH --job-name=dovla_ms_train_debug
+#SBATCH --account=def-yalda_gpu
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=4
+#SBATCH --gres=gpu:nvidia_h100_80gb_hbm3_1g.10gb:1
+#SBATCH --mem=24G
+#SBATCH --time=00:20:00
+#SBATCH --output=outputs/hpc/logs/%x_%j.out
+#SBATCH --error=outputs/hpc/logs/%x_%j.err
+
+set -euo pipefail
+
+PROJECT_DIR="${PROJECT_DIR:-$SLURM_SUBMIT_DIR}"
+DATASET="${DATASET:-$PROJECT_DIR/outputs/hpc/maniskill_debug_cil}"
+RUN_ROOT="${RUN_ROOT:-$PROJECT_DIR/outputs/hpc/maniskill_debug_runs}"
+PYTHON="${PYTHON:-$PROJECT_DIR/.venv/bin/python}"
+export DATASET RUN_ROOT
+
+cd "$PROJECT_DIR"
+mkdir -p outputs/hpc/logs "$RUN_ROOT"
+
+export OMP_NUM_THREADS=1
+export OPENBLAS_NUM_THREADS=1
+export MKL_NUM_THREADS=1
+export DOVLA_TORCH_THREADS=1
+
+test -f "$DATASET/manifest.json"
+"$PYTHON" - <<'PY'
+import json
+import os
+from pathlib import Path
+
+import torch
+
+manifest = json.loads((Path(os.environ["DATASET"]) / "manifest.json").read_text())
+print("torch", torch.__version__, "cuda", torch.cuda.is_available())
+if torch.cuda.is_available():
+    print("gpu", torch.cuda.get_device_name(0))
+print("dataset_records", manifest["record_count"], "groups", manifest["group_count"])
+assert manifest["record_count"] > 0 and manifest["group_count"] >= 4
+PY
+
+COMMON_ARGS=(
+  --dataset "$DATASET"
+  --epochs 10
+  --batch-groups 2
+  --records-per-group 4
+  --pair-count-per-group 6
+  --hidden-dim 128
+  --obs-dim 96
+  --lang-dim 64
+  --action-dim 7
+  --action-horizon 4
+  --effect-dim 16
+  --lr 0.001
+  --device cuda
+  --seed 0
+  --val-fraction 0.25
+)
+
+"$PYTHON" scripts/train_dovla.py \
+  "${COMMON_ARGS[@]}" \
+  --objective lattice_field \
+  --lattice-neighbors 2 \
+  --out "$RUN_ROOT/lattice_field"
+
+"$PYTHON" scripts/train_dovla.py \
+  "${COMMON_ARGS[@]}" \
+  --objective legacy \
+  --out "$RUN_ROOT/legacy"
+
+"$PYTHON" - <<'PY'
+import json
+import os
+from pathlib import Path
+
+root = Path(os.environ["RUN_ROOT"])
+summary = {}
+for name in ("lattice_field", "legacy"):
+    metrics = json.loads((root / name / "metrics.json").read_text())
+    summary[name] = metrics["best"]
+(root / "comparison.json").write_text(json.dumps(summary, indent=2, sort_keys=True) + "\n")
+print(json.dumps(summary, indent=2, sort_keys=True))
+PY

scripts/slurm/train_maniskill_full_array.sbatch

@@ -0,0 +1,74 @@
+#!/bin/bash
+#SBATCH --job-name=dovla_ms_train
+#SBATCH --account=def-yalda_gpu
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=4
+#SBATCH --gres=gpu:nvidia_h100_80gb_hbm3_1g.10gb:1
+#SBATCH --mem=24G
+#SBATCH --time=01:00:00
+#SBATCH --array=0-5%6
+#SBATCH --output=outputs/hpc/logs/%x_%A_%a.out
+#SBATCH --error=outputs/hpc/logs/%x_%A_%a.err
+
+set -euo pipefail
+
+PROJECT_DIR="${PROJECT_DIR:-$SLURM_SUBMIT_DIR}"
+DATASET="${DATASET:-$PROJECT_DIR/outputs/hpc/maniskill_full_k16_n1000_seed0}"
+RUN_ROOT="${RUN_ROOT:-$PROJECT_DIR/outputs/hpc/maniskill_full_runs}"
+PYTHON="${PYTHON:-$PROJECT_DIR/.venv/bin/python}"
+
+TASK_INDEX="${SLURM_ARRAY_TASK_ID:-0}"
+SEED="$((TASK_INDEX / 2))"
+if (( TASK_INDEX % 2 == 0 )); then
+  OBJECTIVE="lattice_field"
+else
+  OBJECTIVE="legacy"
+fi
+OUT_DIR="$RUN_ROOT/$OBJECTIVE/seed_$SEED"
+
+cd "$PROJECT_DIR"
+mkdir -p outputs/hpc/logs "$OUT_DIR"
+
+export OMP_NUM_THREADS=1
+export OPENBLAS_NUM_THREADS=1
+export MKL_NUM_THREADS=1
+export DOVLA_TORCH_THREADS=1
+
+test -f "$DATASET/manifest.json"
+"$PYTHON" - <<PY
+import json
+from pathlib import Path
+import torch
+
+manifest = json.loads((Path("$DATASET") / "manifest.json").read_text())
+assert manifest["group_count"] == 1000
+assert manifest["record_count"] == 16000
+assert torch.cuda.is_available()
+print(
+    "objective=$OBJECTIVE seed=$SEED",
+    "gpu=", torch.cuda.get_device_name(0),
+    "groups=", manifest["group_count"],
+    "records=", manifest["record_count"],
+)
+PY
+
+"$PYTHON" scripts/train_dovla.py \
+  --dataset "$DATASET" \
+  --out "$OUT_DIR" \
+  --epochs 50 \
+  --batch-groups 32 \
+  --records-per-group 16 \
+  --pair-count-per-group 32 \
+  --hidden-dim 256 \
+  --obs-dim 96 \
+  --lang-dim 64 \
+  --action-dim 7 \
+  --action-horizon 4 \
+  --effect-dim 32 \
+  --lr 0.001 \
+  --device cuda \
+  --seed "$SEED" \
+  --val-fraction 0.2 \
+  --objective "$OBJECTIVE" \
+  --lattice-neighbors 32

scripts/slurm/train_maniskill_scaling_array.sbatch

@@ -0,0 +1,82 @@
+#!/bin/bash
+#SBATCH --job-name=dovla_ms_scale_train
+#SBATCH --account=def-yalda_gpu
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=4
+#SBATCH --gres=gpu:nvidia_h100_80gb_hbm3_1g.10gb:1
+#SBATCH --mem=24G
+#SBATCH --time=01:00:00
+#SBATCH --array=0-2%3
+#SBATCH --output=outputs/hpc/logs/%x_%A_%a.out
+#SBATCH --error=outputs/hpc/logs/%x_%A_%a.err
+
+set -euo pipefail
+
+PROJECT_DIR="${PROJECT_DIR:-$SLURM_SUBMIT_DIR}"
+K="${K:?export K before submission}"
+NUM_GROUPS="${NUM_GROUPS:?export NUM_GROUPS before submission}"
+TOTAL_RECORDS="${TOTAL_RECORDS:-16000}"
+DATASET="${DATASET:?export DATASET before submission}"
+RUN_ROOT="${RUN_ROOT:-$PROJECT_DIR/outputs/hpc/maniskill_scaling_runs}"
+PYTHON="${PYTHON:-$PROJECT_DIR/.venv/bin/python}"
+SEED="${SLURM_ARRAY_TASK_ID:-0}"
+BATCH_GROUPS="$((512 / K))"
+PAIR_COUNT="$((2 * K))"
+OUT_DIR="$RUN_ROOT/k_$K/seed_$SEED"
+
+if (( K <= 0 || 512 % K != 0 )); then
+  echo "K must be a positive divisor of 512" >&2
+  exit 2
+fi
+if (( NUM_GROUPS * K != TOTAL_RECORDS )); then
+  echo "fixed-budget invariant failed: NUM_GROUPS*K != TOTAL_RECORDS" >&2
+  exit 2
+fi
+
+cd "$PROJECT_DIR"
+mkdir -p outputs/hpc/logs "$OUT_DIR"
+
+export OMP_NUM_THREADS=1
+export OPENBLAS_NUM_THREADS=1
+export MKL_NUM_THREADS=1
+export DOVLA_TORCH_THREADS=1
+
+test -f "$DATASET/manifest.json"
+"$PYTHON" - <<PY
+import json
+from pathlib import Path
+import torch
+
+manifest = json.loads((Path("$DATASET") / "manifest.json").read_text())
+assert manifest["group_count"] == int("$NUM_GROUPS")
+assert manifest["record_count"] == int("$TOTAL_RECORDS")
+assert manifest["k"] == int("$K")
+assert torch.cuda.is_available()
+print(
+    "K=$K seed=$SEED batch_groups=$BATCH_GROUPS",
+    "gpu=", torch.cuda.get_device_name(0),
+    "groups=", manifest["group_count"],
+    "records=", manifest["record_count"],
+)
+PY
+
+"$PYTHON" scripts/train_dovla.py \
+  --dataset "$DATASET" \
+  --out "$OUT_DIR" \
+  --epochs 50 \
+  --batch-groups "$BATCH_GROUPS" \
+  --records-per-group "$K" \
+  --pair-count-per-group "$PAIR_COUNT" \
+  --hidden-dim 256 \
+  --obs-dim 96 \
+  --lang-dim 64 \
+  --action-dim 7 \
+  --action-horizon 4 \
+  --effect-dim 32 \
+  --lr 0.001 \
+  --device cuda \
+  --seed "$SEED" \
+  --val-fraction 0.2 \
+  --objective lattice_field \
+  --lattice-neighbors 32

scripts/slurm/train_maniskill_visual_array.sbatch

@@ -0,0 +1,79 @@
+#!/bin/bash
+#SBATCH --job-name=dovla_ms_rgb_train
+#SBATCH --account=def-yalda_gpu
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=8
+#SBATCH --gres=gpu:nvidia_h100_80gb_hbm3_1g.10gb:1
+#SBATCH --mem=28G
+#SBATCH --time=03:00:00
+#SBATCH --array=0-2%3
+#SBATCH --output=outputs/hpc/logs/%x_%A_%a.out
+#SBATCH --error=outputs/hpc/logs/%x_%A_%a.err
+
+set -euo pipefail
+
+PROJECT_DIR="${PROJECT_DIR:-$SLURM_SUBMIT_DIR}"
+DATASET="${DATASET:?Set DATASET to a rendered CIL collection}"
+RUN_ROOT="${RUN_ROOT:?Set RUN_ROOT}"
+SCRATCH_ROOT="/scratch/$USER/dovla"
+SIF="${SIF:-$SCRATCH_ROOT/containers/pytorch_2.7.1_cuda12.8.sif}"
+PYTHON="${PYTHON:-$SCRATCH_ROOT/envs/maniskill/bin/python}"
+SEED="${SLURM_ARRAY_TASK_ID:-0}"
+EPOCHS="${EPOCHS:-10}"
+BATCH_GROUPS="${BATCH_GROUPS:-8}"
+HIDDEN_DIM="${HIDDEN_DIM:-256}"
+OUT_DIR="$RUN_ROOT/lattice_field/seed_$SEED"
+
+cd "$PROJECT_DIR"
+mkdir -p "$OUT_DIR"
+module load StdEnv/2023 apptainer/1.4.5
+export OMP_NUM_THREADS=1
+export OPENBLAS_NUM_THREADS=1
+export MKL_NUM_THREADS=1
+export DOVLA_TORCH_THREADS=1
+RUNTIME=(
+  apptainer exec --nv
+  --env "OMP_NUM_THREADS=1,OPENBLAS_NUM_THREADS=1,MKL_NUM_THREADS=1,DOVLA_TORCH_THREADS=1"
+  -B "$PROJECT_DIR:$PROJECT_DIR"
+  -B "/scratch/$USER:/scratch/$USER"
+  "$SIF"
+  "$PYTHON"
+)
+
+"${RUNTIME[@]}" - <<PY
+from dovla_cil.data.datasets import CILDataset
+import torch
+
+dataset = CILDataset("$DATASET")
+assert dataset.group_ids and len(dataset) > 0
+assert all(record.observation_ref for record in dataset.records[: min(256, len(dataset))])
+assert torch.cuda.is_available()
+print(
+    "visual seed=$SEED",
+    "gpu=", torch.cuda.get_device_name(0),
+    "groups=", len(dataset.group_ids),
+    "records=", len(dataset),
+)
+PY
+
+"${RUNTIME[@]}" scripts/train_dovla.py \
+  --dataset "$DATASET" \
+  --out "$OUT_DIR" \
+  --epochs "$EPOCHS" \
+  --batch-groups "$BATCH_GROUPS" \
+  --records-per-group 16 \
+  --pair-count-per-group 32 \
+  --hidden-dim "$HIDDEN_DIM" \
+  --obs-dim 96 \
+  --observation-mode rgb \
+  --lang-dim 64 \
+  --action-dim 8 \
+  --action-horizon 4 \
+  --effect-dim 32 \
+  --lr 0.001 \
+  --device cuda \
+  --seed "$SEED" \
+  --val-fraction 0.2 \
+  --objective lattice_field \
+  --lattice-neighbors 32

scripts/slurm/train_transformer.sbatch

@@ -0,0 +1,72 @@
+#!/bin/bash
+#SBATCH --job-name=dovla_transformer
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=8
+#SBATCH --gres=gpu:1
+#SBATCH --mem=64000M
+#SBATCH --time=48:00:00
+#SBATCH --output=logs/transformer_train_%A_%a.out
+#SBATCH --error=logs/transformer_train_%A_%a.err
+#SBATCH --array=0-2
+
+set -euo pipefail
+
+# DoVLA-Transformer: Pure Transformer Architecture (BREAKTHROUGH)
+# Expected: 42-47% success (vs 38.43% baseline, 36.31% failed Enhanced)
+
+PROJECT_DIR="${PROJECT_DIR:-$PWD}"
+cd "$PROJECT_DIR"
+
+source .venv/bin/activate
+
+DATASET="/scratch/$USER/dovla/experiments/maniskill_presuccess_six_task_collection"
+OUT_DIR="/scratch/$USER/dovla/experiments/cvpr_transformer_model"
+SEED=$SLURM_ARRAY_TASK_ID
+
+mkdir -p "$OUT_DIR/seed_$SEED" logs
+
+echo "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "="
+echo "DoVLA-Transformer: BREAKTHROUGH Architecture"
+echo "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "="
+echo ""
+echo "Pure Transformer Components:"
+echo "  - Multi-head self-attention (8 heads)"
+echo "  - Cross-attention for obs-lang fusion"
+echo "  - 3 Transformer encoder layers"
+echo "  - Positional encoding"
+echo "  - Residual connections everywhere"
+echo ""
+echo "Key Improvements:"
+echo "  - Higher LR: 0.001 (vs 0.0003 failed Enhanced)"
+echo "  - Warmup scheduler: 500 steps"
+echo "  - No custom GNN (proven Transformer only)"
+echo "  - Proper gradient flow (residuals)"
+echo ""
+echo "Dataset: 3,500 groups (fair comparison)"
+echo "Seed: $SEED"
+echo ""
+echo "Expected: 42-47% success"
+echo "vs Baseline: 38.43%"
+echo "vs Enhanced (failed): 36.31%"
+echo ""
+
+python scripts/train_dovla_transformer.py \
+    --dataset "$DATASET" \
+    --out "$OUT_DIR/seed_$SEED" \
+    --d-model 256 \
+    --n-heads 8 \
+    --n-layers 3 \
+    --d-ff 1024 \
+    --epochs 50 \
+    --batch-size 16 \
+    --lr 0.001 \
+    --weight-decay 0.01 \
+    --warmup-steps 500 \
+    --seed $SEED \
+    --device auto
+
+echo ""
+echo "✅ Transformer training complete (seed $SEED)"
+echo ""
+echo "Next: Evaluate and compare"

scripts/slurm/train_transformer_lang.sbatch

@@ -0,0 +1,68 @@
+#!/bin/bash
+#SBATCH --job-name=transformer_lang
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=8
+#SBATCH --gres=gpu:1
+#SBATCH --mem=64000M
+#SBATCH --time=48:00:00
+#SBATCH --output=logs/transformer_lang_%A_%a.out
+#SBATCH --error=logs/transformer_lang_%A_%a.err
+#SBATCH --array=0-2
+
+set -euo pipefail
+
+# DoVLA-Transformer WITH LANGUAGE
+# Expected: 50-55% (from 42-44% baseline)
+# Improvement: +8-11%
+
+PROJECT_DIR="${PROJECT_DIR:-$PWD}"
+cd "$PROJECT_DIR"
+
+source .venv/bin/activate
+
+DATASET="/scratch/$USER/dovla/experiments/maniskill_presuccess_six_task_collection"
+EMBEDDINGS="/scratch/$USER/dovla/experiments/instruction_embeddings.pkl"
+OUT_DIR="/scratch/$USER/dovla/experiments/cvpr_transformer_lang_model"
+SEED=$SLURM_ARRAY_TASK_ID
+
+mkdir -p "$OUT_DIR/seed_$SEED" logs
+
+echo "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "="
+echo "DoVLA-Transformer WITH LANGUAGE"
+echo "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "=" "="
+echo ""
+echo "NEW FEATURE: Instruction embeddings (768-dim)"
+echo "  - Baseline (no language): 42-44%"
+echo "  - WITH language: 50-55% expected"
+echo "  - Improvement: +8-11%"
+echo ""
+echo "Architecture:"
+echo "  - Pure Transformer (8 heads, 3 layers)"
+echo "  - Language dimension: 768"
+echo "  - Cross-attention: obs + lang → context"
+echo ""
+echo "Dataset: 3,500 groups"
+echo "Seed: $SEED"
+echo ""
+
+python scripts/train_transformer_with_language.py \
+    --dataset "$DATASET" \
+    --embeddings "$EMBEDDINGS" \
+    --out "$OUT_DIR/seed_$SEED" \
+    --d-model 256 \
+    --n-heads 8 \
+    --n-layers 3 \
+    --d-ff 1024 \
+    --epochs 50 \
+    --batch-size 16 \
+    --lr 0.001 \
+    --weight-decay 0.01 \
+    --warmup-steps 500 \
+    --seed $SEED \
+    --device auto
+
+echo ""
+echo "✅ Training with language complete (seed $SEED)"
+echo ""
+echo "Next: Evaluate and compare"

scripts/smoke_full_pipeline.py

@@ -0,0 +1,155 @@
+#!/usr/bin/env python
+from __future__ import annotations
+
+import argparse
+import contextlib
+import io
+import os
+import shutil
+import sys
+from pathlib import Path
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+if str(PROJECT_ROOT) not in sys.path:
+    sys.path.insert(0, str(PROJECT_ROOT))
+
+from dovla_cil.eval.causalstress import (  # noqa: E402
+    CausalStressBenchmark,
+    CausalStressConfig,
+    write_metrics_json,
+)
+from dovla_cil.experiments.reports import generate_dataset_report, generate_eval_report
+from dovla_cil.generation.pipeline import generate_cil_dataset, print_generation_summary
+from dovla_cil.tasks.library import ToyTaskLibrary
+from dovla_cil.training.trainer import DoVLATrainer, TrainerConfig
+from dovla_cil.utils.io import ensure_dir, write_jsonl
+from scripts.inspect_shard import main as inspect_main
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(description="Run the full local DoVLA-CIL smoke pipeline.")
+    parser.add_argument("--out", type=Path, default=Path("outputs/smoke_full"))
+    parser.add_argument("--num-tasks", type=int, default=3)
+    parser.add_argument("--states-per-task", type=int, default=4)
+    parser.add_argument("--k", type=int, default=4)
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--shard-size", type=int, default=32)
+    parser.add_argument("--epochs", type=int, default=1)
+    parser.add_argument("--batch-groups", type=int, default=2)
+    parser.add_argument("--records-per-group", type=int, default=4)
+    parser.add_argument("--hidden-dim", type=int, default=64)
+    parser.add_argument("--eval-num-tasks", type=int, default=6)
+    parser.add_argument("--device", default="cpu")
+    parser.add_argument("--no-clean", action="store_true", help="Do not remove an existing output dir.")
+    args = parser.parse_args(argv)
+
+    if args.num_tasks <= 0 or args.states_per_task <= 0 or args.k <= 0:
+        raise ValueError("num-tasks, states-per-task, and k must be positive")
+
+    os.environ.setdefault("OPENCLAUDE_MOCK", "1")
+    if args.out.exists() and not args.no_clean:
+        shutil.rmtree(args.out)
+    output_dir = ensure_dir(args.out)
+    dataset_dir = output_dir / "cil_toy"
+    train_dir = output_dir / "train"
+    report_dir = output_dir / "dataset_report"
+    eval_metrics_path = output_dir / "causalstress" / "metrics.json"
+    eval_report_dir = output_dir / "eval_report"
+    inspect_path = output_dir / "inspect.txt"
+    task_path = output_dir / "tasks.jsonl"
+
+    print("1. Loading built-in toy tasks")
+    tasks = ToyTaskLibrary().list(args.num_tasks)
+    write_jsonl((task.to_dict() for task in tasks), task_path)
+    print(f"   tasks: {task_path}")
+
+    print("2. Generating CIL dataset")
+    generation_summary = generate_cil_dataset(
+        backend="toy",
+        tasks=tasks,
+        out_dir=dataset_dir,
+        num_states_per_task=args.states_per_task,
+        k=args.k,
+        seed=args.seed,
+        shard_size=args.shard_size,
+        inline_observations=True,
+    )
+    print_generation_summary(generation_summary)
+
+    print("3. Inspecting dataset")
+    inspect_output = _capture_stdout(
+        lambda: inspect_main([str(dataset_dir), "--max-rows", str(args.k)])
+    )
+    inspect_path.write_text(inspect_output, encoding="utf-8")
+    print(inspect_output.rstrip())
+
+    print("4. Training DoVLA for one smoke epoch")
+    trainer_result = DoVLATrainer(
+        TrainerConfig(
+            dataset_dir=dataset_dir,
+            output_dir=train_dir,
+            epochs=args.epochs,
+            batch_groups=args.batch_groups,
+            records_per_group=args.records_per_group,
+            pair_count_per_group=args.records_per_group,
+            hidden_dim=args.hidden_dim,
+            learning_rate=1e-3,
+            device=args.device,
+            seed=args.seed,
+            val_fraction=0.25,
+        )
+    ).train()
+    print(f"   checkpoints: {train_dir}")
+    print(f"   best metrics: {trainer_result.get('best', {})}")
+
+    print("5. Evaluating CausalStress")
+    eval_config = CausalStressConfig(
+        backend="toy",
+        num_tasks=args.eval_num_tasks,
+        k=args.k,
+        seed=args.seed,
+    )
+    metrics = CausalStressBenchmark(eval_config).evaluate(train_dir / "best.pt", device=args.device)
+    metrics["config"] = {
+        "backend": "toy",
+        "checkpoint": str(train_dir / "best.pt"),
+        "num_tasks": args.eval_num_tasks,
+        "k": args.k,
+        "seed": args.seed,
+    }
+    write_metrics_json(metrics, eval_metrics_path)
+    print(f"   metrics: {eval_metrics_path}")
+
+    print("6. Writing dataset report")
+    generate_dataset_report(dataset_dir, report_dir, sample_groups=3, seed=args.seed)
+    print(f"   dataset report: {report_dir}")
+
+    print("7. Writing evaluation report")
+    generate_eval_report([eval_metrics_path], eval_report_dir, experiment_name="smoke_full")
+    print(f"   eval report: {eval_report_dir}")
+
+    print("8. Final paths")
+    for label, path in {
+        "root": output_dir,
+        "tasks": task_path,
+        "dataset": dataset_dir,
+        "inspect": inspect_path,
+        "train": train_dir,
+        "checkpoint": train_dir / "best.pt",
+        "causalstress_metrics": eval_metrics_path,
+        "dataset_report": report_dir,
+        "eval_report": eval_report_dir,
+    }.items():
+        print(f"   {label}: {path}")
+    return 0
+
+
+def _capture_stdout(callback) -> str:
+    buffer = io.StringIO()
+    with contextlib.redirect_stdout(buffer):
+        callback()
+    return buffer.getvalue()
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

scripts/smoke_smolvla_checkpoint.py

@@ -0,0 +1,169 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import importlib.metadata
+import json
+import sys
+import time
+from pathlib import Path
+from typing import Any
+
+ROOT = Path(__file__).resolve().parents[1]
+if str(ROOT) not in sys.path:
+    sys.path.insert(0, str(ROOT))
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        description="Load a local SmolVLA checkpoint and write a reproducible smoke manifest."
+    )
+    parser.add_argument("--checkpoint", type=Path, required=True)
+    parser.add_argument(
+        "--vlm-metadata",
+        type=Path,
+        help=(
+            "Local SmolVLM config/tokenizer directory. Required for an offline weight-loading "
+            "smoke test."
+        ),
+    )
+    parser.add_argument("--out", type=Path, required=True)
+    parser.add_argument("--device", default="auto", choices=("auto", "cpu", "cuda"))
+    parser.add_argument(
+        "--metadata-only",
+        action="store_true",
+        help="Validate files and package availability without allocating model weights.",
+    )
+    return parser
+
+
+def smoke_checkpoint(
+    checkpoint: Path,
+    *,
+    device: str = "auto",
+    metadata_only: bool = False,
+    vlm_metadata: Path | None = None,
+) -> dict[str, Any]:
+    checkpoint = checkpoint.expanduser().resolve()
+    required = ("config.json", "model.safetensors")
+    missing = [name for name in required if not (checkpoint / name).is_file()]
+    if missing:
+        raise FileNotFoundError(
+            f"SmolVLA checkpoint is incomplete at {checkpoint}: missing {', '.join(missing)}"
+        )
+
+    result: dict[str, Any] = {
+        "schema_version": "smolvla-checkpoint-smoke/v0",
+        "checkpoint": str(checkpoint),
+        "metadata_only": metadata_only,
+        "required_files": list(required),
+        "package_versions": {
+            name: _package_version(name)
+            for name in ("lerobot", "torch", "transformers", "huggingface-hub")
+        },
+    }
+    if metadata_only:
+        result["ready"] = result["package_versions"]["lerobot"] is not None
+        return result
+
+    if vlm_metadata is None:
+        raise FileNotFoundError(
+            "Offline SmolVLA loading requires --vlm-metadata with local SmolVLM "
+            "config/tokenizer files"
+        )
+    vlm_metadata = vlm_metadata.expanduser().resolve()
+    vlm_required = ("config.json", "preprocessor_config.json", "tokenizer.json")
+    missing_vlm = [name for name in vlm_required if not (vlm_metadata / name).is_file()]
+    if missing_vlm:
+        raise FileNotFoundError(
+            f"SmolVLM metadata is incomplete at {vlm_metadata}: missing {', '.join(missing_vlm)}"
+        )
+
+    try:
+        import torch
+
+        from dovla_cil.eval.smolvla_runtime import (
+            import_smolvla_classes,
+            load_smolvla_config,
+        )
+
+        SmolVLAPolicy, _ = import_smolvla_classes()
+    except ImportError as exc:
+        raise ImportError(
+            'SmolVLA runtime is unavailable. Install isolated dependencies with '
+            '`pip install "lerobot[smolvla]==0.4.3"`. '
+            f"Original import error: {type(exc).__name__}: {exc}"
+        ) from exc
+
+    resolved_device = "cuda" if device == "auto" and torch.cuda.is_available() else device
+    if resolved_device == "auto":
+        resolved_device = "cpu"
+    if resolved_device == "cuda" and not torch.cuda.is_available():
+        raise RuntimeError("CUDA was requested, but torch.cuda.is_available() is false")
+
+    print(json.dumps({"phase": "config_loading"}), flush=True)
+    config = load_smolvla_config(checkpoint, local_files_only=True)
+    config.device = resolved_device
+    config.vlm_model_name = str(vlm_metadata)
+    config.load_vlm_weights = False
+
+    print(json.dumps({"phase": "model_loading", "device": resolved_device}), flush=True)
+    started = time.perf_counter()
+    policy = SmolVLAPolicy.from_pretrained(
+        str(checkpoint),
+        config=config,
+        local_files_only=True,
+    )
+    policy = policy.to(torch.device(resolved_device)).eval()
+    load_seconds = time.perf_counter() - started
+    parameters = sum(parameter.numel() for parameter in policy.parameters())
+    trainable_parameters = sum(
+        parameter.numel() for parameter in policy.parameters() if parameter.requires_grad
+    )
+    print(json.dumps({"phase": "model_ready", "device": resolved_device}), flush=True)
+    result.update(
+        {
+            "ready": True,
+            "device": resolved_device,
+            "vlm_metadata": str(vlm_metadata),
+            "vlm_load_mode": "local_config_then_policy_safetensors",
+            "cuda_device": (
+                torch.cuda.get_device_name(0) if resolved_device == "cuda" else None
+            ),
+            "load_seconds": load_seconds,
+            "parameter_count": parameters,
+            "trainable_parameter_count": trainable_parameters,
+            "policy_class": f"{type(policy).__module__}.{type(policy).__name__}",
+        }
+    )
+    return result
+
+
+def _package_version(name: str) -> str | None:
+    try:
+        return importlib.metadata.version(name)
+    except importlib.metadata.PackageNotFoundError:
+        return None
+
+
+def main() -> int:
+    args = build_parser().parse_args()
+    try:
+        result = smoke_checkpoint(
+            args.checkpoint,
+            device=args.device,
+            metadata_only=args.metadata_only,
+            vlm_metadata=args.vlm_metadata,
+        )
+    except (FileNotFoundError, ImportError, RuntimeError) as exc:
+        print(json.dumps({"ready": False, "error": str(exc)}, indent=2))
+        return 2
+
+    args.out.parent.mkdir(parents=True, exist_ok=True)
+    args.out.write_text(json.dumps(result, indent=2, sort_keys=True), encoding="utf-8")
+    print(json.dumps(result, indent=2, sort_keys=True))
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

scripts/smoke_test.sh

@@ -0,0 +1,27 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+cd "$ROOT_DIR"
+
+export OPENCLAUDE_MOCK="${OPENCLAUDE_MOCK:-1}"
+
+OUT_ROOT="${DOVLA_SMOKE_OUT:-outputs/phase5_smoke}"
+TASKS_PATH="$OUT_ROOT/tasks.jsonl"
+DATASET_DIR="$OUT_ROOT/cil"
+
+mkdir -p "$OUT_ROOT"
+
+python scripts/generate_tasks.py --mock --num-tasks 3 --out "$TASKS_PATH" --seed 0
+python scripts/generate_cil.py \
+  --backend toy \
+  --tasks "$TASKS_PATH" \
+  --out "$DATASET_DIR" \
+  --num-states-per-task 2 \
+  --k 4 \
+  --seed 0 \
+  --shard-size 8 \
+  --inline-observations
+python scripts/inspect_shard.py "$DATASET_DIR"
+
+echo "smoke dataset: $DATASET_DIR"

scripts/train_dovla.py

@@ -0,0 +1,153 @@
+#!/usr/bin/env python
+from __future__ import annotations
+
+import argparse
+import sys
+from dataclasses import fields
+from pathlib import Path
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+if str(PROJECT_ROOT) not in sys.path:
+    sys.path.insert(0, str(PROJECT_ROOT))
+
+from dovla_cil.training.losses import InterventionalLossWeights  # noqa: E402
+from dovla_cil.training.trainer import DoVLATrainer, TrainerConfig  # noqa: E402
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(description="Train the lightweight DoVLA model on CIL data.")
+    parser.add_argument("--dataset", type=Path, required=True)
+    parser.add_argument("--out", type=Path, required=True)
+    parser.add_argument("--epochs", type=int, default=5)
+    parser.add_argument("--batch-groups", type=int, default=8)
+    parser.add_argument("--records-per-group", type=int, default=8)
+    parser.add_argument("--pair-count-per-group", type=int, default=8)
+    parser.add_argument("--hidden-dim", type=int, default=256)
+    parser.add_argument("--obs-dim", type=int, default=32)
+    parser.add_argument(
+        "--observation-mode",
+        choices=("state", "rgb"),
+        default="state",
+        help="Use inline state features or JPEG/HDF5 RGB observation references.",
+    )
+    parser.add_argument("--lang-dim", type=int, default=64)
+    parser.add_argument("--action-dim", type=int, default=8)
+    parser.add_argument("--action-horizon", type=int, default=4)
+    parser.add_argument("--effect-dim", type=int, default=32)
+    parser.add_argument(
+        "--backbone",
+        choices=("native", "clip"),
+        default="native",
+        help="Observation-language backbone. CLIP remains optional and locally loaded.",
+    )
+    parser.add_argument(
+        "--backbone-model",
+        help="Pinned local Hugging Face model directory for the optional CLIP backbone.",
+    )
+    parser.add_argument(
+        "--finetune-backbone",
+        action="store_true",
+        help="Fine-tune pretrained CLIP instead of the default frozen-feature regime.",
+    )
+    parser.add_argument(
+        "--backbone-feature-cache",
+        type=Path,
+        help="Reusable frozen CLIP feature cache shared across seeds.",
+    )
+    parser.add_argument("--backbone-feature-batch-size", type=int, default=64)
+    parser.add_argument("--lr", type=float, default=1e-3)
+    parser.add_argument("--weight-decay", type=float, default=0.0)
+    parser.add_argument("--device", default="auto")
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--val-fraction", type=float, default=0.2)
+    parser.add_argument("--wandb", action="store_true", help="Enable wandb if installed.")
+    parser.add_argument(
+        "--objective",
+        choices=("lattice_field", "legacy"),
+        default="lattice_field",
+        help="Use the proposed interventional field objective or the legacy multi-head ablation.",
+    )
+    parser.add_argument(
+        "--lattice-neighbors",
+        type=int,
+        default=32,
+        help="Nearest action neighbors per node; 32 gives a complete graph for K<=32.",
+    )
+    parser.add_argument(
+        "--pair-scope",
+        choices=("same_state", "cross_state"),
+        default="same_state",
+        help="Choose whether legacy ranking pairs share the exact simulator state.",
+    )
+    parser.add_argument(
+        "--loss-weight",
+        action="append",
+        default=[],
+        metavar="NAME=VALUE",
+        help="Override one loss weight; repeat for multiple weights.",
+    )
+    args = parser.parse_args(argv)
+    try:
+        loss_weights = _parse_loss_weights(args.loss_weight)
+    except ValueError as exc:
+        parser.error(str(exc))
+
+    config = TrainerConfig(
+        dataset_dir=args.dataset,
+        output_dir=args.out,
+        epochs=args.epochs,
+        batch_groups=args.batch_groups,
+        records_per_group=args.records_per_group,
+        pair_count_per_group=args.pair_count_per_group,
+        hidden_dim=args.hidden_dim,
+        obs_dim=args.obs_dim,
+        observation_mode=args.observation_mode,
+        lang_dim=args.lang_dim,
+        action_dim=args.action_dim,
+        action_horizon=args.action_horizon,
+        effect_dim=args.effect_dim,
+        backbone_type=args.backbone,
+        backbone_model=args.backbone_model,
+        backbone_freeze=not args.finetune_backbone,
+        backbone_feature_cache=args.backbone_feature_cache,
+        backbone_feature_batch_size=args.backbone_feature_batch_size,
+        learning_rate=args.lr,
+        weight_decay=args.weight_decay,
+        device=args.device,
+        seed=args.seed,
+        val_fraction=args.val_fraction,
+        wandb=args.wandb,
+        objective=args.objective,
+        lattice_neighbors=args.lattice_neighbors,
+        pair_scope=args.pair_scope,
+        losses=loss_weights,
+    )
+    result = DoVLATrainer(config).train()
+    best = result.get("best", {})
+    print(f"wrote checkpoints to {args.out}")
+    print(f"best val rank_acc={best.get('rank_acc', 0.0):.4f}")
+    return 0
+
+
+def _parse_loss_weights(items: list[str]) -> InterventionalLossWeights:
+    allowed = {field.name for field in fields(InterventionalLossWeights)}
+    values: dict[str, float] = {}
+    for item in items:
+        if "=" not in item:
+            raise ValueError(f"loss weight must use NAME=VALUE syntax: {item!r}")
+        name, raw_value = item.split("=", 1)
+        if name not in allowed:
+            choices = ", ".join(sorted(allowed))
+            raise ValueError(f"unknown loss weight {name!r}; choose one of: {choices}")
+        try:
+            value = float(raw_value)
+        except ValueError as exc:
+            raise ValueError(f"loss weight {name!r} must be numeric") from exc
+        if value < 0:
+            raise ValueError(f"loss weight {name!r} must be non-negative")
+        values[name] = value
+    return InterventionalLossWeights(**values)
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

scripts/train_dovla_attention.py

@@ -0,0 +1,324 @@
+#!/usr/bin/env python
+"""
+Standalone trainer for DoVLA-Attention (CVPR submission)
+
+Single architectural contribution: Transformer attention for action comparison
+- Cross-attention: observation conditions actions
+- Self-attention: models action relationships
+- Pairwise comparison: structured features
+
+Expected: 42-44% success (vs 38.43% MLP baseline)
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import random
+import sys
+from pathlib import Path
+from typing import Optional
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader, Dataset
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+if str(PROJECT_ROOT) not in sys.path:
+    sys.path.insert(0, str(PROJECT_ROOT))
+
+from dovla_cil.models.dovla_attention import DoVLAAttention
+from dovla_cil.data.cil_collection import CILCollection
+
+
+class AttentionTrainingDataset(Dataset):
+    """Dataset for training DoVLA-Attention with pairwise ranking."""
+
+    def __init__(self, collection: CILCollection, group_ids: list[str],
+                 records_per_group: int = 8, pairs_per_group: int = 8):
+        self.collection = collection
+        self.group_ids = group_ids
+        self.records_per_group = records_per_group
+        self.pairs_per_group = pairs_per_group
+
+    def __len__(self):
+        return len(self.group_ids)
+
+    def __getitem__(self, idx):
+        group_id = self.group_ids[idx]
+        records = self.collection.get_group_records(group_id)
+
+        # Sample records
+        if len(records) > self.records_per_group:
+            records = random.sample(records, self.records_per_group)
+
+        # Extract data
+        obs = records[0]["observation"]  # Same state for all
+        actions = [r["action"] for r in records]
+        rewards = [r.get("reward", 0.0) for r in records]
+
+        # Sample pairs for ranking loss
+        pairs = []
+        for _ in range(self.pairs_per_group):
+            i, j = random.sample(range(len(records)), 2)
+            if rewards[i] != rewards[j]:
+                pairs.append((i, j, 1.0 if rewards[i] > rewards[j] else 0.0))
+
+        return {
+            "observation": torch.FloatTensor(obs),
+            "actions": torch.FloatTensor(actions),
+            "rewards": torch.FloatTensor(rewards),
+            "pairs": pairs
+        }
+
+
+def set_seed(seed: int):
+    """Set all random seeds for reproducibility."""
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+def train_epoch(model: nn.Module, dataloader: DataLoader,
+                optimizer: optim.Optimizer, device: torch.device) -> float:
+    """Train for one epoch."""
+    model.train()
+    total_loss = 0.0
+    num_batches = 0
+
+    for batch in dataloader:
+        obs = batch["observation"].to(device)
+        actions = batch["actions"].to(device)
+        rewards = batch["rewards"].to(device)
+
+        optimizer.zero_grad()
+
+        # Forward: get pairwise scores
+        scores = model(obs, actions)  # (batch, K, K)
+
+        # Ranking loss: prefer higher reward actions
+        batch_size, K = rewards.shape
+        loss = 0.0
+        count = 0
+
+        for b in range(batch_size):
+            for i in range(K):
+                for j in range(K):
+                    if i != j and rewards[b, i] != rewards[b, j]:
+                        target = 1.0 if rewards[b, i] > rewards[b, j] else 0.0
+                        pred = torch.sigmoid(scores[b, i, j])
+                        loss += nn.functional.binary_cross_entropy(
+                            pred.unsqueeze(0),
+                            torch.tensor([target], device=device)
+                        )
+                        count += 1
+
+        if count > 0:
+            loss = loss / count
+            loss.backward()
+            optimizer.step()
+
+            total_loss += loss.item()
+            num_batches += 1
+
+    return total_loss / max(num_batches, 1)
+
+
+def evaluate(model: nn.Module, dataloader: DataLoader, device: torch.device) -> dict:
+    """Evaluate model on validation set."""
+    model.eval()
+    correct = 0
+    total = 0
+
+    with torch.no_grad():
+        for batch in dataloader:
+            obs = batch["observation"].to(device)
+            actions = batch["actions"].to(device)
+            rewards = batch["rewards"].to(device)
+
+            scores = model(obs, actions)
+
+            batch_size, K = rewards.shape
+            for b in range(batch_size):
+                for i in range(K):
+                    for j in range(K):
+                        if i != j and rewards[b, i] != rewards[b, j]:
+                            pred = scores[b, i, j] > 0
+                            target = rewards[b, i] > rewards[b, j]
+                            if pred == target:
+                                correct += 1
+                            total += 1
+
+    return {
+        "accuracy": correct / max(total, 1),
+        "correct": correct,
+        "total": total
+    }
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(
+        description="Train DoVLA-Attention for CVPR paper"
+    )
+
+    # Data
+    parser.add_argument("--dataset", type=Path, required=True)
+    parser.add_argument("--out", type=Path, required=True)
+
+    # Architecture
+    parser.add_argument("--hidden-dim", type=int, default=256)
+    parser.add_argument("--n-heads", type=int, default=4)
+    parser.add_argument("--n-layers", type=int, default=2)
+
+    # Training
+    parser.add_argument("--epochs", type=int, default=50)
+    parser.add_argument("--batch-size", type=int, default=16)
+    parser.add_argument("--lr", type=float, default=0.0003)
+    parser.add_argument("--weight-decay", type=float, default=0.01)
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--val-fraction", type=float, default=0.2)
+
+    # System
+    parser.add_argument("--device", default="auto")
+
+    args = parser.parse_args(argv)
+
+    # Setup
+    set_seed(args.seed)
+    args.out.mkdir(parents=True, exist_ok=True)
+
+    if args.device == "auto":
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    else:
+        device = torch.device(args.device)
+
+    print("=" * 70)
+    print("DoVLA-Attention Training (CVPR)")
+    print("=" * 70)
+    print(f"Dataset: {args.dataset}")
+    print(f"Output: {args.out}")
+    print(f"Device: {device}")
+    print(f"Hidden dim: {args.hidden_dim}")
+    print(f"Heads: {args.n_heads}, Layers: {args.n_layers}")
+    print(f"Seed: {args.seed}")
+    print()
+
+    # Load data
+    print("Loading dataset...")
+    collection = CILCollection(args.dataset)
+    all_groups = list(collection.group_ids)
+
+    # Split train/val
+    random.shuffle(all_groups)
+    split_idx = int(len(all_groups) * (1 - args.val_fraction))
+    train_groups = all_groups[:split_idx]
+    val_groups = all_groups[split_idx:]
+
+    print(f"Total groups: {len(all_groups)}")
+    print(f"Train: {len(train_groups)}, Val: {len(val_groups)}")
+    print()
+
+    # Create datasets
+    train_dataset = AttentionTrainingDataset(collection, train_groups)
+    val_dataset = AttentionTrainingDataset(collection, val_groups)
+
+    train_loader = DataLoader(train_dataset, batch_size=args.batch_size,
+                              shuffle=True, num_workers=0)
+    val_loader = DataLoader(val_dataset, batch_size=args.batch_size,
+                           shuffle=False, num_workers=0)
+
+    # Get dims from first sample
+    sample = train_dataset[0]
+    obs_dim = sample["observation"].shape[0]
+    action_dim = sample["actions"].shape[1]
+
+    print(f"Observation dim: {obs_dim}")
+    print(f"Action dim: {action_dim}")
+    print()
+
+    # Create model
+    model = DoVLAAttention(
+        obs_dim=obs_dim,
+        action_dim=action_dim,
+        hidden_dim=args.hidden_dim,
+        n_heads=args.n_heads,
+        n_layers=args.n_layers
+    ).to(device)
+
+    num_params = sum(p.numel() for p in model.parameters())
+    print(f"Model parameters: {num_params:,}")
+    print()
+
+    # Optimizer
+    optimizer = optim.AdamW(model.parameters(), lr=args.lr,
+                           weight_decay=args.weight_decay)
+
+    # Training loop
+    best_acc = 0.0
+    history = []
+
+    print("Starting training...")
+    print()
+
+    for epoch in range(args.epochs):
+        train_loss = train_epoch(model, train_loader, optimizer, device)
+        val_metrics = evaluate(model, val_loader, device)
+
+        val_acc = val_metrics["accuracy"]
+
+        history.append({
+            "epoch": epoch + 1,
+            "train_loss": train_loss,
+            "val_accuracy": val_acc
+        })
+
+        print(f"Epoch {epoch+1:3d}/{args.epochs}: "
+              f"loss={train_loss:.4f}, val_acc={val_acc:.4f}")
+
+        # Save best model
+        if val_acc > best_acc:
+            best_acc = val_acc
+            torch.save({
+                "model_state_dict": model.state_dict(),
+                "epoch": epoch + 1,
+                "val_accuracy": val_acc,
+                "args": vars(args)
+            }, args.out / "best.pt")
+
+    print()
+    print(f"✅ Training complete! Best val accuracy: {best_acc:.4f}")
+
+    # Save training history
+    with open(args.out / "history.json", "w") as f:
+        json.dump(history, f, indent=2)
+
+    # Save final config
+    with open(args.out / "config.json", "w") as f:
+        json.dump({
+            "model": "DoVLA-Attention",
+            "architecture": {
+                "hidden_dim": args.hidden_dim,
+                "n_heads": args.n_heads,
+                "n_layers": args.n_layers
+            },
+            "training": {
+                "epochs": args.epochs,
+                "lr": args.lr,
+                "weight_decay": args.weight_decay,
+                "seed": args.seed
+            },
+            "results": {
+                "best_val_accuracy": best_acc,
+                "num_parameters": num_params
+            }
+        }, f, indent=2)
+
+    print(f"Saved to: {args.out}")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

scripts/train_dovla_enhanced.py

@@ -0,0 +1,407 @@
+#!/usr/bin/env python
+"""
+Enhanced DoVLA-Attention Trainer with SOTA Components
+
+Architecture improvements:
+1. Hierarchical attention (local + global)
+2. Graph neural networks (explicit structure)
+3. Contrastive learning (better embeddings)
+4. Task-adaptive layers (multi-task)
+
+Expected: 44-47% success (vs 38.43% baseline, +5.5-8.5%)
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import random
+import sys
+from pathlib import Path
+from typing import Optional
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader, Dataset
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+if str(PROJECT_ROOT) not in sys.path:
+    sys.path.insert(0, str(PROJECT_ROOT))
+
+from dovla_cil.models.dovla_attention_enhanced import DoVLAAttentionEnhanced
+from dovla_cil.data.datasets import CILDataset
+
+
+class EnhancedTrainingDataset(Dataset):
+    """Dataset for enhanced architecture with task IDs and rewards."""
+
+    def __init__(self, dataset: CILDataset, group_ids: list[str],
+                 records_per_group: int = 16,
+                 max_obs_dim: int = 70, max_act_dim: int = 32):
+        self.dataset = dataset
+        self.group_ids = group_ids
+        self.records_per_group = records_per_group
+        # Pad all observations/actions to fixed max dims for multi-task batching.
+        # This is a standard, fair approach: every method sees the same padded space.
+        self.max_obs_dim = max_obs_dim
+        self.max_act_dim = max_act_dim
+
+        # Task name to ID mapping
+        self.task_map = {
+            "PickCube-v1": 0,
+            "PushCube-v1": 1,
+            "PullCube-v1": 2,
+            "StackCube-v1": 3,
+            "LiftPegUpright-v1": 4,
+            "PegInsertionSide-v1": 5
+        }
+
+    def _pad(self, vec: list[float], target: int) -> list[float]:
+        if len(vec) >= target:
+            return vec[:target]
+        return vec + [0.0] * (target - len(vec))
+
+    def __len__(self):
+        return len(self.group_ids)
+
+    def __getitem__(self, idx):
+        group_id = self.group_ids[idx]
+        records = self.dataset.get_group(group_id)
+
+        # Sample more records for better training
+        if len(records) > self.records_per_group:
+            records = random.sample(records, self.records_per_group)
+
+        # Extract data from CILRecord objects
+        # Observation can be inline dict or reference (use inline if available)
+        obs_data = records[0].observation_inline
+        if obs_data is None or not isinstance(obs_data, dict):
+            raise ValueError(f"No inline observation for group {group_id}")
+
+        # Convert observation dict to flat array
+        if "state" in obs_data:
+            obs = list(obs_data["state"])
+        else:
+            # Flatten all numeric values
+            obs = []
+            for v in obs_data.values():
+                if isinstance(v, list):
+                    obs.extend(v)
+                elif isinstance(v, (int, float)):
+                    obs.append(v)
+
+        # Pad observation to fixed dim for multi-task batching
+        obs = self._pad([float(x) for x in obs], self.max_obs_dim)
+
+        # Extract actions from action_chunk, pad to fixed dim
+        actions = [self._pad(r.action_chunk.flat_values, self.max_act_dim) for r in records]
+
+        # Extract rewards
+        rewards = [r.reward.score for r in records]
+
+        # Get task ID
+        task_name = records[0].task_id
+        task_id = self.task_map.get(task_name, 0)
+
+        return {
+            "observation": torch.FloatTensor(obs),
+            "actions": torch.FloatTensor(actions),
+            "rewards": torch.FloatTensor(rewards),
+            "task_id": torch.LongTensor([task_id])
+        }
+
+
+def collate_fn(batch):
+    """Custom collate to handle variable K and ensure fixed dims."""
+    max_k = max(b["actions"].shape[0] for b in batch)
+
+    batch_size = len(batch)
+    obs_dim = batch[0]["observation"].shape[0]  # Already padded to fixed dim
+    action_dim = batch[0]["actions"].shape[1]   # Already padded to fixed dim
+
+    # Stack observations (all same size after padding)
+    obs_batch = torch.stack([b["observation"] for b in batch])
+
+    # Pad actions to max K in batch
+    actions_batch = torch.zeros(batch_size, max_k, action_dim)
+    rewards_batch = torch.zeros(batch_size, max_k)
+    task_ids = torch.cat([b["task_id"] for b in batch])
+
+    for i, b in enumerate(batch):
+        k = b["actions"].shape[0]
+        actions_batch[i, :k] = b["actions"]
+        rewards_batch[i, :k] = b["rewards"]
+
+    return {
+        "observation": obs_batch,
+        "actions": actions_batch,
+        "rewards": rewards_batch,
+        "task_id": task_ids
+    }
+
+
+def set_seed(seed: int):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+def train_epoch(model: nn.Module, dataloader: DataLoader,
+                optimizer: optim.Optimizer, device: torch.device,
+                contrastive_weight: float = 0.1) -> dict:
+    """Train for one epoch with ranking + contrastive losses."""
+    model.train()
+    total_ranking_loss = 0.0
+    total_contrastive_loss = 0.0
+    num_batches = 0
+
+    for batch in dataloader:
+        obs = batch["observation"].to(device)
+        actions = batch["actions"].to(device)
+        rewards = batch["rewards"].to(device)
+        task_ids = batch["task_id"].to(device).squeeze()
+
+        optimizer.zero_grad()
+
+        # Forward pass
+        scores, contrastive_loss = model(obs, actions, task_ids, rewards)
+
+        # Ranking loss
+        batch_size, K = rewards.shape
+        ranking_loss = 0.0
+        count = 0
+
+        for b in range(batch_size):
+            for i in range(K):
+                for j in range(K):
+                    if i != j and rewards[b, i] != rewards[b, j]:
+                        target = 1.0 if rewards[b, i] > rewards[b, j] else 0.0
+                        pred = torch.sigmoid(scores[b, i, j])
+                        ranking_loss += nn.functional.binary_cross_entropy(
+                            pred.unsqueeze(0),
+                            torch.tensor([target], device=device)
+                        )
+                        count += 1
+
+        if count > 0:
+            ranking_loss = ranking_loss / count
+
+            # Total loss
+            total_loss = ranking_loss
+            if contrastive_loss is not None:
+                total_loss = total_loss + contrastive_weight * contrastive_loss
+
+            total_loss.backward()
+
+            # Gradient clipping for stability
+            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
+
+            optimizer.step()
+
+            total_ranking_loss += ranking_loss.item()
+            if contrastive_loss is not None:
+                total_contrastive_loss += contrastive_loss.item()
+            num_batches += 1
+
+    return {
+        "ranking_loss": total_ranking_loss / max(num_batches, 1),
+        "contrastive_loss": total_contrastive_loss / max(num_batches, 1)
+    }
+
+
+def evaluate(model: nn.Module, dataloader: DataLoader, device: torch.device) -> dict:
+    """Evaluate model."""
+    model.eval()
+    correct = 0
+    total = 0
+
+    with torch.no_grad():
+        for batch in dataloader:
+            obs = batch["observation"].to(device)
+            actions = batch["actions"].to(device)
+            rewards = batch["rewards"].to(device)
+            task_ids = batch["task_id"].to(device).squeeze()
+
+            scores, _ = model(obs, actions, task_ids, None)
+
+            batch_size, K = rewards.shape
+            for b in range(batch_size):
+                for i in range(K):
+                    for j in range(K):
+                        if i != j and rewards[b, i] != rewards[b, j]:
+                            pred = scores[b, i, j] > 0
+                            target = rewards[b, i] > rewards[b, j]
+                            if pred == target:
+                                correct += 1
+                            total += 1
+
+    return {
+        "accuracy": correct / max(total, 1)
+    }
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(
+        description="Train Enhanced DoVLA-Attention for CVPR"
+    )
+
+    # Data
+    parser.add_argument("--dataset", type=Path, required=True)
+    parser.add_argument("--out", type=Path, required=True)
+
+    # Architecture
+    parser.add_argument("--hidden-dim", type=int, default=256)
+    parser.add_argument("--n-heads", type=int, default=4)
+    parser.add_argument("--n-layers", type=int, default=3)
+
+    # Training
+    parser.add_argument("--epochs", type=int, default=50)
+    parser.add_argument("--batch-size", type=int, default=16)
+    parser.add_argument("--lr", type=float, default=0.0003)
+    parser.add_argument("--weight-decay", type=float, default=0.01)
+    parser.add_argument("--contrastive-weight", type=float, default=0.1)
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--val-fraction", type=float, default=0.2)
+
+    # System
+    parser.add_argument("--device", default="auto")
+
+    args = parser.parse_args(argv)
+
+    set_seed(args.seed)
+    args.out.mkdir(parents=True, exist_ok=True)
+
+    if args.device == "auto":
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    else:
+        device = torch.device(args.device)
+
+    print("=" * 70)
+    print("Enhanced DoVLA-Attention Training (CVPR)")
+    print("=" * 70)
+    print(f"Dataset: {args.dataset}")
+    print(f"Device: {device}")
+    print(f"Architecture: Hierarchical + Graph + Contrastive + Task-Adaptive")
+    print(f"Hidden: {args.hidden_dim}, Heads: {args.n_heads}, Layers: {args.n_layers}")
+    print(f"Seed: {args.seed}")
+    print()
+
+    # Load data
+    print("Loading dataset...")
+    dataset = CILDataset(args.dataset)
+    all_groups = list(dataset.group_ids)
+
+    random.shuffle(all_groups)
+    split_idx = int(len(all_groups) * (1 - args.val_fraction))
+    train_groups = all_groups[:split_idx]
+    val_groups = all_groups[split_idx:]
+
+    print(f"Total: {len(all_groups)}, Train: {len(train_groups)}, Val: {len(val_groups)}")
+    print()
+
+    # Datasets
+    train_dataset = EnhancedTrainingDataset(dataset, train_groups, records_per_group=16)
+    val_dataset = EnhancedTrainingDataset(dataset, val_groups, records_per_group=16)
+
+    train_loader = DataLoader(train_dataset, batch_size=args.batch_size,
+                              shuffle=True, num_workers=0, collate_fn=collate_fn)
+    val_loader = DataLoader(val_dataset, batch_size=args.batch_size,
+                           shuffle=False, num_workers=0, collate_fn=collate_fn)
+
+    # Get dimensions
+    sample = train_dataset[0]
+    obs_dim = sample["observation"].shape[0]
+    action_dim = sample["actions"].shape[1]
+
+    print(f"Observation dim: {obs_dim}, Action dim: {action_dim}")
+    print()
+
+    # Create enhanced model
+    model = DoVLAAttentionEnhanced(
+        obs_dim=obs_dim,
+        action_dim=action_dim,
+        hidden_dim=args.hidden_dim,
+        n_heads=args.n_heads,
+        n_layers=args.n_layers,
+        num_tasks=6,
+        use_contrastive=True,
+        use_graph=True,
+        use_task_adaptive=True
+    ).to(device)
+
+    num_params = sum(p.numel() for p in model.parameters())
+    print(f"Model parameters: {num_params:,}")
+    print()
+
+    # Optimizer
+    optimizer = optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
+
+    # Learning rate scheduler
+    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs)
+
+    # Training
+    best_acc = 0.0
+    history = []
+
+    print("Starting training...")
+    print()
+
+    for epoch in range(args.epochs):
+        train_metrics = train_epoch(model, train_loader, optimizer, device, args.contrastive_weight)
+        val_metrics = evaluate(model, val_loader, device)
+        scheduler.step()
+
+        val_acc = val_metrics["accuracy"]
+
+        history.append({
+            "epoch": epoch + 1,
+            "ranking_loss": train_metrics["ranking_loss"],
+            "contrastive_loss": train_metrics["contrastive_loss"],
+            "val_accuracy": val_acc,
+            "lr": scheduler.get_last_lr()[0]
+        })
+
+        print(f"Epoch {epoch+1:3d}/{args.epochs}: "
+              f"rank_loss={train_metrics['ranking_loss']:.4f}, "
+              f"contr_loss={train_metrics['contrastive_loss']:.4f}, "
+              f"val_acc={val_acc:.4f}")
+
+        if val_acc > best_acc:
+            best_acc = val_acc
+            torch.save({
+                "model_state_dict": model.state_dict(),
+                "epoch": epoch + 1,
+                "val_accuracy": val_acc,
+                "args": vars(args)
+            }, args.out / "best.pt")
+
+    print()
+    print(f"✅ Training complete! Best val accuracy: {best_acc:.4f}")
+
+    # Save
+    with open(args.out / "history.json", "w") as f:
+        json.dump(history, f, indent=2)
+
+    with open(args.out / "config.json", "w") as f:
+        json.dump({
+            "model": "DoVLA-Attention-Enhanced",
+            "components": ["hierarchical_attention", "graph_nn", "contrastive", "task_adaptive"],
+            "architecture": {
+                "hidden_dim": args.hidden_dim,
+                "n_heads": args.n_heads,
+                "n_layers": args.n_layers
+            },
+            "results": {
+                "best_val_accuracy": best_acc,
+                "num_parameters": num_params
+            }
+        }, f, indent=2)
+
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

scripts/train_dovla_transformer.py

@@ -0,0 +1,366 @@
+#!/usr/bin/env python
+"""
+Train DoVLA-Transformer with proper Transformer training recipe.
+
+Key improvements over failed Enhanced:
+1. Higher learning rate (0.001 vs 0.0003)
+2. Warmup scheduler (standard for Transformer)
+3. Gradient clipping: 1.0 (standard)
+4. Pure ranking loss (no contrastive)
+5. Standard Transformer components (proven)
+
+Expected: 42-47% success
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import random
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader, Dataset
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+if str(PROJECT_ROOT) not in sys.path:
+    sys.path.insert(0, str(PROJECT_ROOT))
+
+from dovla_cil.models.dovla_transformer import DoVLATransformer
+from dovla_cil.data.datasets import CILDataset
+
+
+class TransformerTrainingDataset(Dataset):
+    """Dataset for DoVLA-Transformer training."""
+
+    def __init__(self, dataset: CILDataset, group_ids: list[str],
+                 records_per_group: int = 16, max_obs_dim: int = 70, max_act_dim: int = 32):
+        self.dataset = dataset
+        self.group_ids = group_ids
+        self.records_per_group = records_per_group
+        self.max_obs_dim = max_obs_dim
+        self.max_act_dim = max_act_dim
+
+    def _pad(self, vec: list[float], target: int) -> list[float]:
+        if len(vec) >= target:
+            return vec[:target]
+        return vec + [0.0] * (target - len(vec))
+
+    def __len__(self):
+        return len(self.group_ids)
+
+    def __getitem__(self, idx):
+        group_id = self.group_ids[idx]
+        records = self.dataset.get_group(group_id)
+
+        if len(records) > self.records_per_group:
+            records = random.sample(records, self.records_per_group)
+
+        # Observation
+        obs_data = records[0].observation_inline
+        if "state" in obs_data:
+            obs = list(obs_data["state"])
+        else:
+            obs = []
+            for v in obs_data.values():
+                if isinstance(v, list):
+                    obs.extend(v)
+                elif isinstance(v, (int, float)):
+                    obs.append(v)
+        obs = self._pad([float(x) for x in obs], self.max_obs_dim)
+
+        # Actions
+        actions = [self._pad(r.action_chunk.flat_values, self.max_act_dim) for r in records]
+
+        # Rewards
+        rewards = [r.reward.score for r in records]
+
+        return {
+            "observation": torch.FloatTensor(obs),
+            "actions": torch.FloatTensor(actions),
+            "rewards": torch.FloatTensor(rewards)
+        }
+
+
+def collate_fn(batch):
+    """Collate with padding to max K in batch."""
+    max_k = max(b["actions"].shape[0] for b in batch)
+    batch_size = len(batch)
+    obs_dim = batch[0]["observation"].shape[0]
+    action_dim = batch[0]["actions"].shape[1]
+
+    obs_batch = torch.stack([b["observation"] for b in batch])
+    actions_batch = torch.zeros(batch_size, max_k, action_dim)
+    rewards_batch = torch.zeros(batch_size, max_k)
+
+    for i, b in enumerate(batch):
+        k = b["actions"].shape[0]
+        actions_batch[i, :k] = b["actions"]
+        rewards_batch[i, :k] = b["rewards"]
+
+    return {
+        "observation": obs_batch,
+        "actions": actions_batch,
+        "rewards": rewards_batch
+    }
+
+
+def set_seed(seed: int):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+def get_cosine_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps):
+    """Cosine schedule with linear warmup (standard for Transformer)."""
+    def lr_lambda(current_step):
+        if current_step < num_warmup_steps:
+            return float(current_step) / float(max(1, num_warmup_steps))
+        progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))
+        return max(0.0, 0.5 * (1.0 + np.cos(np.pi * progress)))
+
+    return optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
+
+
+def train_epoch(model: nn.Module, dataloader: DataLoader,
+                optimizer: optim.Optimizer, device: torch.device) -> float:
+    """Train for one epoch."""
+    model.train()
+    total_loss = 0.0
+    num_batches = 0
+
+    for batch in dataloader:
+        obs = batch["observation"].to(device)
+        actions = batch["actions"].to(device)
+        rewards = batch["rewards"].to(device)
+
+        optimizer.zero_grad()
+
+        scores = model(obs, actions)
+
+        # Ranking loss
+        batch_size, K = rewards.shape
+        loss = 0.0
+        count = 0
+
+        for b in range(batch_size):
+            for i in range(K):
+                for j in range(K):
+                    if i != j and rewards[b, i] != rewards[b, j]:
+                        target = 1.0 if rewards[b, i] > rewards[b, j] else 0.0
+                        pred = torch.sigmoid(scores[b, i, j])
+                        loss += nn.functional.binary_cross_entropy(
+                            pred.unsqueeze(0),
+                            torch.tensor([target], device=device)
+                        )
+                        count += 1
+
+        if count > 0:
+            loss = loss / count
+            loss.backward()
+
+            # Gradient clipping
+            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
+
+            optimizer.step()
+
+            total_loss += loss.item()
+            num_batches += 1
+
+    return total_loss / max(num_batches, 1)
+
+
+def evaluate(model: nn.Module, dataloader: DataLoader, device: torch.device) -> dict:
+    """Evaluate model - proper action selection accuracy."""
+    model.eval()
+    correct_top1 = 0
+    total_groups = 0
+
+    with torch.no_grad():
+        for batch in dataloader:
+            obs = batch["observation"].to(device)
+            actions = batch["actions"].to(device)
+            rewards = batch["rewards"].to(device)
+
+            scores_matrix = model(obs, actions)
+
+            # Aggregate pairwise to per-action scores
+            batch_size, K = rewards.shape
+            for b in range(batch_size):
+                # Sum wins for each action
+                action_scores = scores_matrix[b].sum(dim=1).cpu().tolist()
+
+                # Select best
+                selected = max(range(K), key=lambda i: action_scores[i])
+                best_utility = max(rewards[b].cpu().tolist())
+
+                if abs(rewards[b, selected].item() - best_utility) < 1e-6:
+                    correct_top1 += 1
+                total_groups += 1
+
+    return {
+        "top1_accuracy": correct_top1 / max(total_groups, 1)
+    }
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(description="Train DoVLA-Transformer")
+
+    # Data
+    parser.add_argument("--dataset", type=Path, required=True)
+    parser.add_argument("--out", type=Path, required=True)
+
+    # Architecture
+    parser.add_argument("--d-model", type=int, default=256)
+    parser.add_argument("--n-heads", type=int, default=8)
+    parser.add_argument("--n-layers", type=int, default=3)
+    parser.add_argument("--d-ff", type=int, default=1024)
+
+    # Training
+    parser.add_argument("--epochs", type=int, default=50)
+    parser.add_argument("--batch-size", type=int, default=16)
+    parser.add_argument("--lr", type=float, default=0.001)  # Higher than failed Enhanced
+    parser.add_argument("--weight-decay", type=float, default=0.01)
+    parser.add_argument("--warmup-steps", type=int, default=500)
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--val-fraction", type=float, default=0.2)
+
+    # System
+    parser.add_argument("--device", default="auto")
+
+    args = parser.parse_args(argv)
+
+    set_seed(args.seed)
+    args.out.mkdir(parents=True, exist_ok=True)
+
+    if args.device == "auto":
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    else:
+        device = torch.device(args.device)
+
+    print("=" * 70)
+    print("DoVLA-Transformer Training (BREAKTHROUGH)")
+    print("=" * 70)
+    print(f"Dataset: {args.dataset}")
+    print(f"Device: {device}")
+    print(f"Architecture: Pure Transformer (d={args.d_model}, heads={args.n_heads}, layers={args.n_layers})")
+    print(f"LR: {args.lr} (higher than failed Enhanced 0.0003)")
+    print(f"Warmup: {args.warmup_steps} steps")
+    print(f"Seed: {args.seed}")
+    print()
+
+    # Load data
+    print("Loading dataset...")
+    dataset = CILDataset(args.dataset)
+    all_groups = list(dataset.group_ids)
+
+    random.shuffle(all_groups)
+    split_idx = int(len(all_groups) * (1 - args.val_fraction))
+    train_groups = all_groups[:split_idx]
+    val_groups = all_groups[split_idx:]
+
+    print(f"Total: {len(all_groups)}, Train: {len(train_groups)}, Val: {len(val_groups)}")
+    print()
+
+    # Datasets
+    train_dataset = TransformerTrainingDataset(dataset, train_groups)
+    val_dataset = TransformerTrainingDataset(dataset, val_groups)
+
+    train_loader = DataLoader(train_dataset, batch_size=args.batch_size,
+                              shuffle=True, num_workers=0, collate_fn=collate_fn)
+    val_loader = DataLoader(val_dataset, batch_size=args.batch_size,
+                           shuffle=False, num_workers=0, collate_fn=collate_fn)
+
+    # Model
+    model = DoVLATransformer(
+        obs_dim=70,
+        action_dim=32,
+        lang_dim=0,
+        d_model=args.d_model,
+        n_heads=args.n_heads,
+        n_layers=args.n_layers,
+        d_ff=args.d_ff,
+        dropout=0.1
+    ).to(device)
+
+    num_params = sum(p.numel() for p in model.parameters())
+    print(f"Model parameters: {num_params:,}")
+    print()
+
+    # Optimizer & Scheduler
+    optimizer = optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
+
+    num_training_steps = len(train_loader) * args.epochs
+    scheduler = get_cosine_schedule_with_warmup(optimizer, args.warmup_steps, num_training_steps)
+
+    # Training
+    best_acc = 0.0
+    history = []
+
+    print("Starting training...")
+    print()
+
+    for epoch in range(args.epochs):
+        train_loss = train_epoch(model, train_loader, optimizer, device)
+        val_metrics = evaluate(model, val_loader, device)
+        scheduler.step()
+
+        val_acc = val_metrics["top1_accuracy"]
+
+        history.append({
+            "epoch": epoch + 1,
+            "train_loss": train_loss,
+            "val_top1_accuracy": val_acc,
+            "lr": scheduler.get_last_lr()[0]
+        })
+
+        print(f"Epoch {epoch+1:3d}/{args.epochs}: "
+              f"loss={train_loss:.4f}, val_top1={val_acc:.4f}, lr={scheduler.get_last_lr()[0]:.6f}")
+
+        if val_acc > best_acc:
+            best_acc = val_acc
+            torch.save({
+                "model_state_dict": model.state_dict(),
+                "epoch": epoch + 1,
+                "val_top1_accuracy": val_acc,
+                "args": vars(args)
+            }, args.out / "best.pt")
+
+    print()
+    print(f"✅ Training complete! Best val top-1: {best_acc:.4f}")
+
+    # Save
+    with open(args.out / "history.json", "w") as f:
+        json.dump(history, f, indent=2)
+
+    with open(args.out / "config.json", "w") as f:
+        json.dump({
+            "model": "DoVLA-Transformer",
+            "architecture": {
+                "d_model": args.d_model,
+                "n_heads": args.n_heads,
+                "n_layers": args.n_layers,
+                "d_ff": args.d_ff
+            },
+            "training": {
+                "lr": args.lr,
+                "warmup_steps": args.warmup_steps,
+                "epochs": args.epochs,
+                "seed": args.seed
+            },
+            "results": {
+                "best_val_top1_accuracy": best_acc,
+                "num_parameters": num_params
+            }
+        }, f, indent=2)
+
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

scripts/train_hybrid_direct.py

@@ -0,0 +1,348 @@
+#!/usr/bin/env python
+"""
+Train DoVLA-Hybrid with DIRECT scoring (NOT pairwise).
+
+Key improvement: Predict reward + success directly
+Expected: 45-48% baseline (vs 37% pairwise)
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import random
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torch.utils.data import DataLoader, Dataset
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+if str(PROJECT_ROOT) not in sys.path:
+    sys.path.insert(0, str(PROJECT_ROOT))
+
+from dovla_cil.models.dovla_hybrid import DoVLAHybrid
+from dovla_cil.data.datasets import CILDataset
+
+
+class HybridTrainingDataset(Dataset):
+    """Dataset for hybrid direct scoring."""
+
+    def __init__(self, dataset: CILDataset, group_ids: list[str],
+                 records_per_group: int = 16, max_obs_dim: int = 70, max_act_dim: int = 32):
+        self.dataset = dataset
+        self.group_ids = group_ids
+        self.records_per_group = records_per_group
+        self.max_obs_dim = max_obs_dim
+        self.max_act_dim = max_act_dim
+
+    def _pad(self, vec: list[float], target: int) -> list[float]:
+        if len(vec) >= target:
+            return vec[:target]
+        return vec + [0.0] * (target - len(vec))
+
+    def __len__(self):
+        return len(self.group_ids)
+
+    def __getitem__(self, idx):
+        group_id = self.group_ids[idx]
+        records = self.dataset.get_group(group_id)
+
+        if len(records) > self.records_per_group:
+            records = random.sample(records, self.records_per_group)
+
+        # Observation
+        obs_data = records[0].observation_inline
+        if "state" in obs_data:
+            obs = list(obs_data["state"])
+        else:
+            obs = []
+            for v in obs_data.values():
+                if isinstance(v, list):
+                    obs.extend(v)
+                elif isinstance(v, (int, float)):
+                    obs.append(v)
+        obs = self._pad([float(x) for x in obs], self.max_obs_dim)
+
+        # Actions
+        actions = [self._pad(r.action_chunk.flat_values, self.max_act_dim) for r in records]
+
+        # Rewards (direct targets!)
+        rewards = [r.reward.score for r in records]
+
+        # Success labels (direct targets!)
+        successes = [float(r.reward.terminal_success) for r in records]
+
+        return {
+            "observation": torch.FloatTensor(obs),
+            "actions": torch.FloatTensor(actions),
+            "rewards": torch.FloatTensor(rewards),
+            "successes": torch.FloatTensor(successes)
+        }
+
+
+def collate_fn(batch):
+    """Collate with padding."""
+    max_k = max(b["actions"].shape[0] for b in batch)
+    batch_size = len(batch)
+    obs_dim = batch[0]["observation"].shape[0]
+    action_dim = batch[0]["actions"].shape[1]
+
+    obs_batch = torch.stack([b["observation"] for b in batch])
+    actions_batch = torch.zeros(batch_size, max_k, action_dim)
+    rewards_batch = torch.zeros(batch_size, max_k)
+    successes_batch = torch.zeros(batch_size, max_k)
+
+    for i, b in enumerate(batch):
+        k = b["actions"].shape[0]
+        actions_batch[i, :k] = b["actions"]
+        rewards_batch[i, :k] = b["rewards"]
+        successes_batch[i, :k] = b["successes"]
+
+    return {
+        "observation": obs_batch,
+        "actions": actions_batch,
+        "rewards": rewards_batch,
+        "successes": successes_batch
+    }
+
+
+def set_seed(seed: int):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+def get_cosine_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps):
+    def lr_lambda(current_step):
+        if current_step < num_warmup_steps:
+            return float(current_step) / float(max(1, num_warmup_steps))
+        progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))
+        return max(0.0, 0.5 * (1.0 + np.cos(np.pi * progress)))
+    return optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
+
+
+def train_epoch(model: nn.Module, dataloader: DataLoader,
+                optimizer: optim.Optimizer, device: torch.device) -> dict:
+    """Train for one epoch with DIRECT scoring."""
+    model.train()
+    total_reward_loss = 0.0
+    total_success_loss = 0.0
+    num_batches = 0
+
+    for batch in dataloader:
+        obs = batch["observation"].to(device)
+        actions = batch["actions"].to(device)
+        target_rewards = batch["rewards"].to(device)
+        target_successes = batch["successes"].to(device)
+
+        optimizer.zero_grad()
+
+        # Forward: predict rewards and success probs DIRECTLY
+        pred_rewards, pred_success_probs = model(obs, actions)
+
+        # Loss 1: MSE for reward prediction
+        reward_loss = F.mse_loss(pred_rewards, target_rewards)
+
+        # Loss 2: BCE for success prediction
+        success_loss = F.binary_cross_entropy(pred_success_probs, target_successes)
+
+        # Combined loss
+        loss = reward_loss + success_loss
+
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
+        optimizer.step()
+
+        total_reward_loss += reward_loss.item()
+        total_success_loss += success_loss.item()
+        num_batches += 1
+
+    return {
+        "reward_loss": total_reward_loss / max(num_batches, 1),
+        "success_loss": total_success_loss / max(num_batches, 1),
+        "total_loss": (total_reward_loss + total_success_loss) / max(num_batches, 1)
+    }
+
+
+def evaluate(model: nn.Module, dataloader: DataLoader, device: torch.device) -> dict:
+    """Evaluate with DIRECT selection."""
+    model.eval()
+    correct_top1 = 0
+    total_groups = 0
+
+    with torch.no_grad():
+        for batch in dataloader:
+            obs = batch["observation"].to(device)
+            actions = batch["actions"].to(device)
+            target_rewards = batch["rewards"].to(device)
+
+            # Predict directly
+            pred_rewards, pred_success_probs = model(obs, actions)
+
+            # Hybrid selection: success_prob * predicted_reward
+            hybrid_scores = pred_success_probs * pred_rewards
+
+            batch_size, K = target_rewards.shape
+            for b in range(batch_size):
+                # Select action with highest hybrid score
+                selected = hybrid_scores[b].argmax().item()
+
+                # Check if selected best reward
+                best_reward = target_rewards[b].max().item()
+                if abs(target_rewards[b, selected].item() - best_reward) < 1e-6:
+                    correct_top1 += 1
+                total_groups += 1
+
+    return {"top1_accuracy": correct_top1 / max(total_groups, 1)}
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(description="Train DoVLA-Hybrid (Direct Scoring)")
+
+    parser.add_argument("--dataset", type=Path, required=True)
+    parser.add_argument("--out", type=Path, required=True)
+    parser.add_argument("--d-model", type=int, default=256)
+    parser.add_argument("--n-heads", type=int, default=8)
+    parser.add_argument("--n-layers", type=int, default=3)
+    parser.add_argument("--d-ff", type=int, default=1024)
+    parser.add_argument("--epochs", type=int, default=50)
+    parser.add_argument("--batch-size", type=int, default=16)
+    parser.add_argument("--lr", type=float, default=0.001)
+    parser.add_argument("--weight-decay", type=float, default=0.01)
+    parser.add_argument("--warmup-steps", type=int, default=500)
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--val-fraction", type=float, default=0.2)
+    parser.add_argument("--device", default="auto")
+
+    args = parser.parse_args(argv)
+    set_seed(args.seed)
+    args.out.mkdir(parents=True, exist_ok=True)
+
+    if args.device == "auto":
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    else:
+        device = torch.device(args.device)
+
+    print("=" * 70)
+    print("DoVLA-Hybrid: DIRECT Scoring (NOT Pairwise)")
+    print("=" * 70)
+    print(f"Dataset: {args.dataset}")
+    print(f"Device: {device}")
+    print(f"Approach: Predict reward + success DIRECTLY")
+    print(f"Expected: 45-48% (vs 37% pairwise baseline)")
+    print()
+
+    # Load data
+    dataset = CILDataset(args.dataset)
+    all_groups = list(dataset.group_ids)
+    random.shuffle(all_groups)
+    split_idx = int(len(all_groups) * (1 - args.val_fraction))
+    train_groups = all_groups[:split_idx]
+    val_groups = all_groups[split_idx:]
+
+    print(f"Total: {len(all_groups)}, Train: {len(train_groups)}, Val: {len(val_groups)}")
+    print()
+
+    train_dataset = HybridTrainingDataset(dataset, train_groups)
+    val_dataset = HybridTrainingDataset(dataset, val_groups)
+
+    train_loader = DataLoader(train_dataset, batch_size=args.batch_size,
+                              shuffle=True, num_workers=0, collate_fn=collate_fn)
+    val_loader = DataLoader(val_dataset, batch_size=args.batch_size,
+                           shuffle=False, num_workers=0, collate_fn=collate_fn)
+
+    # Model
+    model = DoVLAHybrid(
+        obs_dim=70,
+        action_dim=32,
+        lang_dim=0,
+        d_model=args.d_model,
+        n_heads=args.n_heads,
+        n_layers=args.n_layers,
+        d_ff=args.d_ff,
+        dropout=0.1
+    ).to(device)
+
+    num_params = sum(p.numel() for p in model.parameters())
+    print(f"Model parameters: {num_params:,}")
+    print()
+
+    optimizer = optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
+    num_training_steps = len(train_loader) * args.epochs
+    scheduler = get_cosine_schedule_with_warmup(optimizer, args.warmup_steps, num_training_steps)
+
+    best_acc = 0.0
+    history = []
+
+    print("Starting training...")
+    print()
+
+    for epoch in range(args.epochs):
+        train_metrics = train_epoch(model, train_loader, optimizer, device)
+        val_metrics = evaluate(model, val_loader, device)
+        scheduler.step()
+
+        val_acc = val_metrics["top1_accuracy"]
+
+        history.append({
+            "epoch": epoch + 1,
+            "train_reward_loss": train_metrics["reward_loss"],
+            "train_success_loss": train_metrics["success_loss"],
+            "train_total_loss": train_metrics["total_loss"],
+            "val_top1_accuracy": val_acc,
+            "lr": scheduler.get_last_lr()[0]
+        })
+
+        print(f"Epoch {epoch+1:3d}/{args.epochs}: "
+              f"r_loss={train_metrics['reward_loss']:.4f}, "
+              f"s_loss={train_metrics['success_loss']:.4f}, "
+              f"val_top1={val_acc:.4f}")
+
+        if val_acc > best_acc:
+            best_acc = val_acc
+            torch.save({
+                "model_state_dict": model.state_dict(),
+                "epoch": epoch + 1,
+                "val_top1_accuracy": val_acc,
+                "args": vars(args)
+            }, args.out / "best.pt")
+
+    print()
+    print(f"✅ Training complete! Best val top-1: {best_acc:.4f}")
+
+    with open(args.out / "history.json", "w") as f:
+        json.dump(history, f, indent=2)
+
+    with open(args.out / "config.json", "w") as f:
+        json.dump({
+            "model": "DoVLA-Hybrid-Direct",
+            "approach": "direct_scoring",
+            "architecture": {
+                "d_model": args.d_model,
+                "n_heads": args.n_heads,
+                "n_layers": args.n_layers,
+                "d_ff": args.d_ff
+            },
+            "training": {
+                "lr": args.lr,
+                "warmup_steps": args.warmup_steps,
+                "epochs": args.epochs,
+                "seed": args.seed
+            },
+            "results": {
+                "best_val_top1_accuracy": best_acc,
+                "num_parameters": num_params
+            }
+        }, f, indent=2)
+
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

scripts/train_transformer_with_language.py

@@ -0,0 +1,368 @@
+#!/usr/bin/env python
+"""
+Train DoVLA-Transformer WITH LANGUAGE EMBEDDINGS.
+
+This is the improved version that uses instruction embeddings.
+Expected improvement: +8-11% (50-55% from 42-44% baseline)
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import pickle
+import random
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader, Dataset
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+if str(PROJECT_ROOT) not in sys.path:
+    sys.path.insert(0, str(PROJECT_ROOT))
+
+from dovla_cil.models.dovla_transformer import DoVLATransformer
+from dovla_cil.data.datasets import CILDataset
+
+
+class TransformerLangDataset(Dataset):
+    """Dataset for DoVLA-Transformer with language embeddings."""
+
+    def __init__(self, dataset: CILDataset, group_ids: list[str],
+                 embeddings: dict, records_per_group: int = 16,
+                 max_obs_dim: int = 70, max_act_dim: int = 32):
+        self.dataset = dataset
+        self.group_ids = group_ids
+        self.embeddings = embeddings  # {group_id: embedding}
+        self.records_per_group = records_per_group
+        self.max_obs_dim = max_obs_dim
+        self.max_act_dim = max_act_dim
+
+    def _pad(self, vec: list[float], target: int) -> list[float]:
+        if len(vec) >= target:
+            return vec[:target]
+        return vec + [0.0] * (target - len(vec))
+
+    def __len__(self):
+        return len(self.group_ids)
+
+    def __getitem__(self, idx):
+        group_id = self.group_ids[idx]
+        records = self.dataset.get_group(group_id)
+
+        if len(records) > self.records_per_group:
+            records = random.sample(records, self.records_per_group)
+
+        # Observation
+        obs_data = records[0].observation_inline
+        if "state" in obs_data:
+            obs = list(obs_data["state"])
+        else:
+            obs = []
+            for v in obs_data.values():
+                if isinstance(v, list):
+                    obs.extend(v)
+                elif isinstance(v, (int, float)):
+                    obs.append(v)
+        obs = self._pad([float(x) for x in obs], self.max_obs_dim)
+
+        # Actions
+        actions = [self._pad(r.action_chunk.flat_values, self.max_act_dim) for r in records]
+
+        # Rewards
+        rewards = [r.reward.score for r in records]
+
+        # Language embedding
+        lang_emb = self.embeddings.get(group_id, np.zeros(768))
+
+        return {
+            "observation": torch.FloatTensor(obs),
+            "actions": torch.FloatTensor(actions),
+            "rewards": torch.FloatTensor(rewards),
+            "language": torch.FloatTensor(lang_emb)  # NEW!
+        }
+
+
+def collate_fn(batch):
+    """Collate with padding to max K in batch."""
+    max_k = max(b["actions"].shape[0] for b in batch)
+    batch_size = len(batch)
+    obs_dim = batch[0]["observation"].shape[0]
+    action_dim = batch[0]["actions"].shape[1]
+    lang_dim = batch[0]["language"].shape[0]
+
+    obs_batch = torch.stack([b["observation"] for b in batch])
+    actions_batch = torch.zeros(batch_size, max_k, action_dim)
+    rewards_batch = torch.zeros(batch_size, max_k)
+    lang_batch = torch.stack([b["language"] for b in batch])  # NEW!
+
+    for i, b in enumerate(batch):
+        k = b["actions"].shape[0]
+        actions_batch[i, :k] = b["actions"]
+        rewards_batch[i, :k] = b["rewards"]
+
+    return {
+        "observation": obs_batch,
+        "actions": actions_batch,
+        "rewards": rewards_batch,
+        "language": lang_batch  # NEW!
+    }
+
+
+def set_seed(seed: int):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+def get_cosine_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps):
+    """Cosine schedule with linear warmup."""
+    def lr_lambda(current_step):
+        if current_step < num_warmup_steps:
+            return float(current_step) / float(max(1, num_warmup_steps))
+        progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))
+        return max(0.0, 0.5 * (1.0 + np.cos(np.pi * progress)))
+
+    return optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
+
+
+def train_epoch(model: nn.Module, dataloader: DataLoader,
+                optimizer: optim.Optimizer, device: torch.device) -> float:
+    """Train for one epoch."""
+    model.train()
+    total_loss = 0.0
+    num_batches = 0
+
+    for batch in dataloader:
+        obs = batch["observation"].to(device)
+        actions = batch["actions"].to(device)
+        rewards = batch["rewards"].to(device)
+        lang = batch["language"].to(device)  # NEW!
+
+        optimizer.zero_grad()
+
+        scores = model(obs, actions, lang)  # Pass language!
+
+        # Ranking loss
+        batch_size, K = rewards.shape
+        loss = 0.0
+        count = 0
+
+        for b in range(batch_size):
+            for i in range(K):
+                for j in range(K):
+                    if i != j and rewards[b, i] != rewards[b, j]:
+                        target = 1.0 if rewards[b, i] > rewards[b, j] else 0.0
+                        pred = torch.sigmoid(scores[b, i, j])
+                        loss += nn.functional.binary_cross_entropy(
+                            pred.unsqueeze(0),
+                            torch.tensor([target], device=device)
+                        )
+                        count += 1
+
+        if count > 0:
+            loss = loss / count
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
+            optimizer.step()
+            total_loss += loss.item()
+            num_batches += 1
+
+    return total_loss / max(num_batches, 1)
+
+
+def evaluate(model: nn.Module, dataloader: DataLoader, device: torch.device) -> dict:
+    """Evaluate model."""
+    model.eval()
+    correct_top1 = 0
+    total_groups = 0
+
+    with torch.no_grad():
+        for batch in dataloader:
+            obs = batch["observation"].to(device)
+            actions = batch["actions"].to(device)
+            rewards = batch["rewards"].to(device)
+            lang = batch["language"].to(device)  # NEW!
+
+            scores_matrix = model(obs, actions, lang)  # Pass language!
+
+            batch_size, K = rewards.shape
+            for b in range(batch_size):
+                action_scores = scores_matrix[b].sum(dim=1).cpu().tolist()
+                selected = max(range(K), key=lambda i: action_scores[i])
+                best_utility = max(rewards[b].cpu().tolist())
+
+                if abs(rewards[b, selected].item() - best_utility) < 1e-6:
+                    correct_top1 += 1
+                total_groups += 1
+
+    return {"top1_accuracy": correct_top1 / max(total_groups, 1)}
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(description="Train DoVLA-Transformer with Language")
+
+    # Data
+    parser.add_argument("--dataset", type=Path, required=True)
+    parser.add_argument("--embeddings", type=Path, required=True)  # NEW!
+    parser.add_argument("--out", type=Path, required=True)
+
+    # Architecture
+    parser.add_argument("--d-model", type=int, default=256)
+    parser.add_argument("--n-heads", type=int, default=8)
+    parser.add_argument("--n-layers", type=int, default=3)
+    parser.add_argument("--d-ff", type=int, default=1024)
+
+    # Training
+    parser.add_argument("--epochs", type=int, default=50)
+    parser.add_argument("--batch-size", type=int, default=16)
+    parser.add_argument("--lr", type=float, default=0.001)
+    parser.add_argument("--weight-decay", type=float, default=0.01)
+    parser.add_argument("--warmup-steps", type=int, default=500)
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--val-fraction", type=float, default=0.2)
+    parser.add_argument("--device", default="auto")
+
+    args = parser.parse_args(argv)
+
+    set_seed(args.seed)
+    args.out.mkdir(parents=True, exist_ok=True)
+
+    if args.device == "auto":
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    else:
+        device = torch.device(args.device)
+
+    print("=" * 70)
+    print("DoVLA-Transformer Training WITH LANGUAGE")
+    print("=" * 70)
+    print(f"Dataset: {args.dataset}")
+    print(f"Embeddings: {args.embeddings}")
+    print(f"Device: {device}")
+    print(f"Language dimension: 768")
+    print(f"Expected improvement: +8-11% (to 50-55%)")
+    print(f"Seed: {args.seed}")
+    print()
+
+    # Load embeddings
+    print("Loading instruction embeddings...")
+    with open(args.embeddings, 'rb') as f:
+        embeddings = pickle.load(f)
+    print(f"Loaded {len(embeddings)} embeddings")
+    print()
+
+    # Load data
+    print("Loading dataset...")
+    dataset = CILDataset(args.dataset)
+    all_groups = list(dataset.group_ids)
+
+    random.shuffle(all_groups)
+    split_idx = int(len(all_groups) * (1 - args.val_fraction))
+    train_groups = all_groups[:split_idx]
+    val_groups = all_groups[split_idx:]
+
+    print(f"Total: {len(all_groups)}, Train: {len(train_groups)}, Val: {len(val_groups)}")
+    print()
+
+    # Datasets
+    train_dataset = TransformerLangDataset(dataset, train_groups, embeddings)
+    val_dataset = TransformerLangDataset(dataset, val_groups, embeddings)
+
+    train_loader = DataLoader(train_dataset, batch_size=args.batch_size,
+                              shuffle=True, num_workers=0, collate_fn=collate_fn)
+    val_loader = DataLoader(val_dataset, batch_size=args.batch_size,
+                           shuffle=False, num_workers=0, collate_fn=collate_fn)
+
+    # Model
+    model = DoVLATransformer(
+        obs_dim=70,
+        action_dim=32,
+        lang_dim=768,  # Enable language!
+        d_model=args.d_model,
+        n_heads=args.n_heads,
+        n_layers=args.n_layers,
+        d_ff=args.d_ff,
+        dropout=0.1
+    ).to(device)
+
+    num_params = sum(p.numel() for p in model.parameters())
+    print(f"Model parameters: {num_params:,}")
+    print()
+
+    # Optimizer & Scheduler
+    optimizer = optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
+    num_training_steps = len(train_loader) * args.epochs
+    scheduler = get_cosine_schedule_with_warmup(optimizer, args.warmup_steps, num_training_steps)
+
+    # Training
+    best_acc = 0.0
+    history = []
+
+    print("Starting training...")
+    print()
+
+    for epoch in range(args.epochs):
+        train_loss = train_epoch(model, train_loader, optimizer, device)
+        val_metrics = evaluate(model, val_loader, device)
+        scheduler.step()
+
+        val_acc = val_metrics["top1_accuracy"]
+
+        history.append({
+            "epoch": epoch + 1,
+            "train_loss": train_loss,
+            "val_top1_accuracy": val_acc,
+            "lr": scheduler.get_last_lr()[0]
+        })
+
+        print(f"Epoch {epoch+1:3d}/{args.epochs}: "
+              f"loss={train_loss:.4f}, val_top1={val_acc:.4f}, lr={scheduler.get_last_lr()[0]:.6f}")
+
+        if val_acc > best_acc:
+            best_acc = val_acc
+            torch.save({
+                "model_state_dict": model.state_dict(),
+                "epoch": epoch + 1,
+                "val_top1_accuracy": val_acc,
+                "args": vars(args)
+            }, args.out / "best.pt")
+
+    print()
+    print(f"✅ Training complete! Best val top-1: {best_acc:.4f}")
+
+    # Save
+    with open(args.out / "history.json", "w") as f:
+        json.dump(history, f, indent=2)
+
+    with open(args.out / "config.json", "w") as f:
+        json.dump({
+            "model": "DoVLA-Transformer-Language",
+            "language_dim": 768,
+            "architecture": {
+                "d_model": args.d_model,
+                "n_heads": args.n_heads,
+                "n_layers": args.n_layers,
+                "d_ff": args.d_ff
+            },
+            "training": {
+                "lr": args.lr,
+                "warmup_steps": args.warmup_steps,
+                "epochs": args.epochs,
+                "seed": args.seed
+            },
+            "results": {
+                "best_val_top1_accuracy": best_acc,
+                "num_parameters": num_params
+            }
+        }, f, indent=2)
+
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())