diff --git a/scripts/slurm/generate_6task_h16.sbatch b/scripts/slurm/generate_6task_h16.sbatch new file mode 100644 index 0000000000000000000000000000000000000000..eda70bc77dc2fe595d3575f0e83af08702c13ae4 --- /dev/null +++ b/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" diff --git a/scripts/slurm/generate_cil_array.sbatch b/scripts/slurm/generate_cil_array.sbatch new file mode 100644 index 0000000000000000000000000000000000000000..29ed98c7af925e47d06b07de173cc9d5e2d85aea --- /dev/null +++ b/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:-}" +# 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[@]}" diff --git a/scripts/slurm/generate_embeddings.sbatch b/scripts/slurm/generate_embeddings.sbatch new file mode 100644 index 0000000000000000000000000000000000000000..b948fb37eb6f6207d35bde217ec281b4d5d11b58 --- /dev/null +++ b/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" diff --git a/scripts/slurm/horizon_sweep_pickcube.sbatch b/scripts/slurm/horizon_sweep_pickcube.sbatch new file mode 100644 index 0000000000000000000000000000000000000000..35ef8939986c0c7980de3067b8a038ddab972973 --- /dev/null +++ b/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 diff --git a/scripts/slurm/install_smolvla_env.sbatch b/scripts/slurm/install_smolvla_env.sbatch new file mode 100644 index 0000000000000000000000000000000000000000..38c7635ca3527d5c3ce8f1a02cd32a3928495310 --- /dev/null +++ b/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')]" diff --git a/scripts/slurm/make_maniskill_collection.sbatch b/scripts/slurm/make_maniskill_collection.sbatch new file mode 100644 index 0000000000000000000000000000000000000000..eb969a75163829cfe4e004d1acb5a78c44c7dfeb --- /dev/null +++ b/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" diff --git a/scripts/slurm/maniskill_lattice_debug.sbatch b/scripts/slurm/maniskill_lattice_debug.sbatch new file mode 100644 index 0000000000000000000000000000000000000000..d93d61effefd7591bf9bac549b2944e44969c320 --- /dev/null +++ b/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" diff --git a/scripts/slurm/maniskill_lattice_full.sbatch b/scripts/slurm/maniskill_lattice_full.sbatch new file mode 100644 index 0000000000000000000000000000000000000000..9ef9ed67101c45f6c977f343884c55361950102e --- /dev/null +++ b/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" diff --git a/scripts/slurm/maniskill_multitask_pilot.sbatch b/scripts/slurm/maniskill_multitask_pilot.sbatch new file mode 100644 index 0000000000000000000000000000000000000000..f8b50948faa344861b9c85a59e59bafcf6796c80 --- /dev/null +++ b/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" diff --git a/scripts/slurm/phase_a1_generate_10k.sbatch b/scripts/slurm/phase_a1_generate_10k.sbatch new file mode 100644 index 0000000000000000000000000000000000000000..65c9eed90626d8aefde6221826657a3a5fbc826c --- /dev/null +++ b/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" diff --git a/scripts/slurm/phase_a1_generate_10k_enhanced.sbatch b/scripts/slurm/phase_a1_generate_10k_enhanced.sbatch new file mode 100644 index 0000000000000000000000000000000000000000..72050885eb83c39a7faecf1f344e5dda227fc76d --- /dev/null +++ b/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" diff --git a/scripts/slurm/phase_a1_revised_enhanced.sbatch b/scripts/slurm/phase_a1_revised_enhanced.sbatch new file mode 100644 index 0000000000000000000000000000000000000000..b8d77bc9df575769acc0f4bdd58240bbc9c12f55 --- /dev/null +++ b/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" diff --git a/scripts/slurm/phase_a1b_train_enhanced.sbatch b/scripts/slurm/phase_a1b_train_enhanced.sbatch new file mode 100644 index 0000000000000000000000000000000000000000..380d7dc5d09a373b4a9ff4e73504522d70111c75 --- /dev/null +++ b/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" diff --git a/scripts/slurm/phase_a2_train_large_model.sbatch b/scripts/slurm/phase_a2_train_large_model.sbatch new file mode 100644 index 0000000000000000000000000000000000000000..478f3bdaf6dea8ad16f516107ed1ca1a9ab0b81e --- /dev/null +++ b/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" diff --git a/scripts/slurm/phase_a3_eval_large_model.sbatch b/scripts/slurm/phase_a3_eval_large_model.sbatch new file mode 100644 index 0000000000000000000000000000000000000000..e15b7db881263442b94233205f9b3d67ef38a586 --- /dev/null +++ b/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)" diff --git a/scripts/slurm/phase_a4_hparam_sweep.sbatch b/scripts/slurm/phase_a4_hparam_sweep.sbatch new file mode 100644 index 0000000000000000000000000000000000000000..322e01c4990b587d16bd808298347b2b3e5508a4 --- /dev/null +++ b/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" diff --git a/scripts/slurm/phase_a5_horizon_sweep.sbatch b/scripts/slurm/phase_a5_horizon_sweep.sbatch new file mode 100644 index 0000000000000000000000000000000000000000..37b9b4f7b1a678d073e0bbc5353403f911fc119a --- /dev/null +++ b/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" diff --git a/scripts/slurm/phase_b_generate_12tasks.sbatch b/scripts/slurm/phase_b_generate_12tasks.sbatch new file mode 100644 index 0000000000000000000000000000000000000000..bf546225b971a68e8fd7b0b4ce6cb4940d397513 --- /dev/null +++ b/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" diff --git a/scripts/slurm/phase_b_train_12tasks.sbatch b/scripts/slurm/phase_b_train_12tasks.sbatch new file mode 100644 index 0000000000000000000000000000000000000000..4981f6c0136fc22bc5620fc687ad29539479ee94 --- /dev/null +++ b/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" diff --git a/scripts/slurm/plan_c_generate_10k.sbatch b/scripts/slurm/plan_c_generate_10k.sbatch new file mode 100644 index 0000000000000000000000000000000000000000..be70c1bc0b75ab6118bdde58f6c4593802bcf43e --- /dev/null +++ b/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" diff --git a/scripts/slurm/prepare_maniskill_baselines.sbatch b/scripts/slurm/prepare_maniskill_baselines.sbatch new file mode 100644 index 0000000000000000000000000000000000000000..ecf79ed0366809d28573e5f70367c33e177b0215 --- /dev/null +++ b/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 diff --git a/scripts/slurm/render_maniskill_multitask.sbatch b/scripts/slurm/render_maniskill_multitask.sbatch new file mode 100644 index 0000000000000000000000000000000000000000..eaa369234b8f643b5d655eb9edd157729a26dd71 --- /dev/null +++ b/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" diff --git a/scripts/slurm/render_maniskill_observations.sbatch b/scripts/slurm/render_maniskill_observations.sbatch new file mode 100644 index 0000000000000000000000000000000000000000..83423c03c21e99c4bed83710f37ce9665d8f211e --- /dev/null +++ b/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" diff --git a/scripts/slurm/run_external_vla_baseline.sbatch b/scripts/slurm/run_external_vla_baseline.sbatch new file mode 100644 index 0000000000000000000000000000000000000000..2ad92f4088d5091627c4d8ec28a1c24b5f598da7 --- /dev/null +++ b/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[@]}" diff --git a/scripts/slurm/run_scaling.sbatch b/scripts/slurm/run_scaling.sbatch new file mode 100644 index 0000000000000000000000000000000000000000..f0b6e5892a17ed4ae1fb9b329fa12f784a3b3b96 --- /dev/null +++ b/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" diff --git a/scripts/slurm/run_smolvla_cil_baseline.sbatch b/scripts/slurm/run_smolvla_cil_baseline.sbatch new file mode 100644 index 0000000000000000000000000000000000000000..7d03d64fd5ac7c1111ce0dfcf79aa799e5bf0947 --- /dev/null +++ b/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 diff --git a/scripts/slurm/smoke_smolvla_checkpoint.sbatch b/scripts/slurm/smoke_smolvla_checkpoint.sbatch new file mode 100644 index 0000000000000000000000000000000000000000..996fd55f6e19ab675e7b9273aa7be3342b131551 --- /dev/null +++ b/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 diff --git a/scripts/slurm/train_attention_model.sbatch b/scripts/slurm/train_attention_model.sbatch new file mode 100644 index 0000000000000000000000000000000000000000..bdfa82690cdb3a45e1d7435cf7b4ee1d96b6d6ae --- /dev/null +++ b/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%)" diff --git a/scripts/slurm/train_dovla.sbatch b/scripts/slurm/train_dovla.sbatch new file mode 100644 index 0000000000000000000000000000000000000000..d5e48bcffdef0fbaf35e205cda0b45f6d1afd763 --- /dev/null +++ b/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" diff --git a/scripts/slurm/train_enhanced_model.sbatch b/scripts/slurm/train_enhanced_model.sbatch new file mode 100644 index 0000000000000000000000000000000000000000..c2fa01a51cadc973edb18672740e4a968a50c2a0 --- /dev/null +++ b/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" diff --git a/scripts/slurm/train_h16_policy.sbatch b/scripts/slurm/train_h16_policy.sbatch new file mode 100644 index 0000000000000000000000000000000000000000..793ae7a7516966dd849a8b93a5ad88f4e297dfe3 --- /dev/null +++ b/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" diff --git a/scripts/slurm/train_hybrid_direct.sbatch b/scripts/slurm/train_hybrid_direct.sbatch new file mode 100644 index 0000000000000000000000000000000000000000..5df652a79bfee2faaffce0aa600b368da0fe8b20 --- /dev/null +++ b/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 "" diff --git a/scripts/slurm/train_maniskill_baseline_array.sbatch b/scripts/slurm/train_maniskill_baseline_array.sbatch new file mode 100644 index 0000000000000000000000000000000000000000..710ae294338d4316a256263578933e418588e837 --- /dev/null +++ b/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[@]}" diff --git a/scripts/slurm/train_maniskill_collection_array.sbatch b/scripts/slurm/train_maniskill_collection_array.sbatch new file mode 100644 index 0000000000000000000000000000000000000000..2d17d6703c0475a2349b605299f2c14c0d617c14 --- /dev/null +++ b/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[@]}" - < 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 diff --git a/scripts/slurm/train_maniskill_full_array.sbatch b/scripts/slurm/train_maniskill_full_array.sbatch new file mode 100644 index 0000000000000000000000000000000000000000..660ca8f4a03ac490884e9a62d18d7a33c33d0f61 --- /dev/null +++ b/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" - <&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" - < 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 diff --git a/scripts/slurm/train_transformer.sbatch b/scripts/slurm/train_transformer.sbatch new file mode 100644 index 0000000000000000000000000000000000000000..6c0e7b543fe55f56d457fd75186349dd5da87516 --- /dev/null +++ b/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" diff --git a/scripts/slurm/train_transformer_lang.sbatch b/scripts/slurm/train_transformer_lang.sbatch new file mode 100644 index 0000000000000000000000000000000000000000..8d54d3c80076be39c18fe571ad6c2bfb46781177 --- /dev/null +++ b/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" diff --git a/scripts/smoke_full_pipeline.py b/scripts/smoke_full_pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..641fca4eefe37a0a0f91a34c492b5033dffcc489 --- /dev/null +++ b/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()) diff --git a/scripts/smoke_smolvla_checkpoint.py b/scripts/smoke_smolvla_checkpoint.py new file mode 100644 index 0000000000000000000000000000000000000000..2b0de0d8e4e8f7ad360fceb6f19ada89345071f3 --- /dev/null +++ b/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()) diff --git a/scripts/smoke_test.sh b/scripts/smoke_test.sh new file mode 100644 index 0000000000000000000000000000000000000000..c476e2ef5bf76d5b17e1078c552a11663839ffbc --- /dev/null +++ b/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" diff --git a/scripts/train_dovla.py b/scripts/train_dovla.py new file mode 100644 index 0000000000000000000000000000000000000000..afb635ab568ea92d7ea89098d5f23d8369fd21db --- /dev/null +++ b/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()) diff --git a/scripts/train_dovla_attention.py b/scripts/train_dovla_attention.py new file mode 100644 index 0000000000000000000000000000000000000000..6927255e07b4d03b3a946e6bd461df8281803a50 --- /dev/null +++ b/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()) diff --git a/scripts/train_dovla_enhanced.py b/scripts/train_dovla_enhanced.py new file mode 100644 index 0000000000000000000000000000000000000000..951644eb693c20d518022d7cfe1e41f2fa409775 --- /dev/null +++ b/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()) diff --git a/scripts/train_dovla_transformer.py b/scripts/train_dovla_transformer.py new file mode 100644 index 0000000000000000000000000000000000000000..6b63da83a773317a6132cf1b7e89cb1754765d72 --- /dev/null +++ b/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()) diff --git a/scripts/train_hybrid_direct.py b/scripts/train_hybrid_direct.py new file mode 100644 index 0000000000000000000000000000000000000000..751fae254ce1d93ed7268674b7591ff563c8680c --- /dev/null +++ b/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()) diff --git a/scripts/train_transformer_with_language.py b/scripts/train_transformer_with_language.py new file mode 100644 index 0000000000000000000000000000000000000000..1efa76f0f77a843927dc96b2df0fccd5aa3423bb --- /dev/null +++ b/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()) diff --git a/scripts/verify_external_checkpoint.py b/scripts/verify_external_checkpoint.py new file mode 100644 index 0000000000000000000000000000000000000000..738a3a8a41feeba9aed5356dfd2e462a2b853a04 --- /dev/null +++ b/scripts/verify_external_checkpoint.py @@ -0,0 +1,129 @@ +#!/usr/bin/env python +from __future__ import annotations + +import argparse +import hashlib +import json +from pathlib import Path + +DEFAULT_REQUIRED = { + "smolvla": [ + "config.json", + "model.safetensors", + "policy_preprocessor.json", + "policy_postprocessor.json", + ], + "clip": [ + "config.json", + "pytorch_model.bin", + "preprocessor_config.json", + "tokenizer_config.json", + ], +} + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Verify a staged external checkpoint and write a SHA256 manifest." + ) + parser.add_argument("--checkpoint", required=True, help="Local checkpoint directory.") + parser.add_argument( + "--out", + default=None, + help="Manifest path. Defaults inside checkpoint dir.", + ) + parser.add_argument("--model-family", default="smolvla", choices=sorted(DEFAULT_REQUIRED)) + parser.add_argument("--repo-id", default=None) + parser.add_argument("--revision", default=None) + parser.add_argument( + "--required-file", + action="append", + default=None, + help="Required file relative to checkpoint dir. May be passed multiple times.", + ) + return parser.parse_args() + + +def main() -> int: + args = parse_args() + checkpoint = Path(args.checkpoint) + required = args.required_file or DEFAULT_REQUIRED[args.model_family] + manifest = build_manifest( + checkpoint, + model_family=args.model_family, + repo_id=args.repo_id, + revision=args.revision, + required_files=required, + ) + out = Path(args.out) if args.out else checkpoint / "dovla_download_manifest.json" + out.parent.mkdir(parents=True, exist_ok=True) + out.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n", encoding="utf-8") + print(json.dumps(_summary(manifest), indent=2, sort_keys=True)) + return 0 if manifest["ready"] else 2 + + +def build_manifest( + checkpoint: Path, + *, + model_family: str, + repo_id: str | None = None, + revision: str | None = None, + required_files: list[str], +) -> dict[str, object]: + if not checkpoint.exists(): + files: list[dict[str, object]] = [] + else: + files = [ + _file_row(path, checkpoint) + for path in sorted(checkpoint.rglob("*")) + if path.is_file() + ] + present = {str(row["path"]) for row in files} + missing = [name for name in required_files if name not in present] + return { + "schema_version": "dovla-external-checkpoint-manifest/v0", + "model_family": model_family, + "repo_id": repo_id, + "revision": revision, + "local_dir": str(checkpoint.resolve()), + "ready": not missing, + "required_files": list(required_files), + "missing_required_files": missing, + "file_count": len(files), + "total_bytes": sum(int(row["size_bytes"]) for row in files), + "files": files, + } + + +def _file_row(path: Path, root: Path) -> dict[str, object]: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return { + "path": str(path.relative_to(root)), + "size_bytes": path.stat().st_size, + "sha256": digest.hexdigest(), + } + + +def _summary(manifest: dict[str, object]) -> dict[str, object]: + files = manifest["files"] + assert isinstance(files, list) + largest = max( + (row for row in files if isinstance(row, dict)), + key=lambda row: int(row["size_bytes"]), + default=None, + ) + return { + "ready": manifest["ready"], + "model_family": manifest["model_family"], + "file_count": manifest["file_count"], + "total_bytes": manifest["total_bytes"], + "missing_required_files": manifest["missing_required_files"], + "largest_file": largest, + } + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000000000000000000000000000000000000..a07525d0bb180fca7f904caddb5404c787ea3421 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,8 @@ +from __future__ import annotations + +import os + +# Login nodes and small CI runners often impose strict process/thread limits. Keep numerical tests +# deterministic and leave production training thread counts to the launcher configuration. +for variable in ("OMP_NUM_THREADS", "OPENBLAS_NUM_THREADS", "MKL_NUM_THREADS"): + os.environ.setdefault(variable, "1") diff --git a/tests/test_baselines.py b/tests/test_baselines.py new file mode 100644 index 0000000000000000000000000000000000000000..59be413e15732f38582741349fbce18b841a9730 --- /dev/null +++ b/tests/test_baselines.py @@ -0,0 +1,114 @@ +from __future__ import annotations + +import subprocess +import sys +from pathlib import Path + +from dovla_cil.data.datasets import CILDataset +from dovla_cil.experiments.baselines import ( + BaselineConfig, + loss_weights_for_baseline, + prepare_dataset_for_baseline, +) +from dovla_cil.generation.pipeline import generate_cil_dataset +from dovla_cil.tasks.library import built_in_toy_tasks +from dovla_cil.utils.io import read_json + + +def _make_dataset(tmp_path: Path) -> Path: + dataset_dir = tmp_path / "cil" + generate_cil_dataset( + backend="toy", + tasks=built_in_toy_tasks()[:2], + out_dir=dataset_dir, + num_states_per_task=1, + k=4, + seed=9, + shard_size=8, + inline_observations=True, + ) + return dataset_dir + + +def test_expert_only_dataset_has_one_record_per_group(tmp_path: Path) -> None: + dataset_dir = _make_dataset(tmp_path) + prepared = prepare_dataset_for_baseline( + dataset_dir, "expert_only_bc", tmp_path / "expert_only" + ) + dataset = CILDataset(prepared) + assert len(dataset.records) == len(dataset.group_ids) + assert all(len(dataset.get_group(group_id)) == 1 for group_id in dataset.group_ids) + + +def test_random_negatives_mode_can_generate(tmp_path: Path) -> None: + dataset_dir = _make_dataset(tmp_path) + prepared = prepare_dataset_for_baseline( + dataset_dir, "random_negatives", tmp_path / "random_negatives" + ) + dataset = CILDataset(prepared) + assert any(record.candidate_type == "random_negative" for record in dataset.records) + assert (prepared / "baseline_metadata.json").exists() + + +def test_world_model_auxiliary_sets_loss_weights() -> None: + weights = loss_weights_for_baseline("world_model_auxiliary") + assert weights.weight("effect") == 1.0 + assert weights.weight("progress") == 1.0 + assert weights.weight("rank") == 0.0 + assert weights.weight("regret") == 0.0 + + +def test_cross_state_baseline_is_measured_not_placeholder(tmp_path: Path) -> None: + dataset_dir = _make_dataset(tmp_path) + prepared = prepare_dataset_for_baseline( + dataset_dir, "cross_state_negatives", tmp_path / "cross_state" + ) + + metadata = read_json(prepared / "baseline_metadata.json") + assert metadata["approximate"] is False + + +def test_baseline_config_model_dump(tmp_path: Path) -> None: + config = BaselineConfig( + baseline="expert_only_bc", + dataset=tmp_path / "dataset", + out=tmp_path / "out", + ) + payload = config.model_dump() + assert payload["baseline"] == "expert_only_bc" + + +def test_baseline_cli_smoke_runs(tmp_path: Path) -> None: + dataset_dir = _make_dataset(tmp_path) + out_dir = tmp_path / "run" + subprocess.run( + [ + sys.executable, + "scripts/run_baseline.py", + "--baseline", + "expert_only_bc", + "--dataset", + str(dataset_dir), + "--out", + str(out_dir), + "--epochs", + "1", + "--batch-groups", + "1", + "--records-per-group", + "1", + "--hidden-dim", + "32", + "--eval-num-tasks", + "2", + "--eval-k", + "4", + ], + check=True, + capture_output=True, + text=True, + ) + assert (out_dir / "train" / "best.pt").exists() + metrics = read_json(out_dir / "metrics.json") + assert metrics["baseline"] == "expert_only_bc" + assert "pairwise_ranking_accuracy" in metrics["eval"] diff --git a/tests/test_causalstress.py b/tests/test_causalstress.py new file mode 100644 index 0000000000000000000000000000000000000000..c34d7af2a9b110598d0271f6fd0eff039f2bf2fb --- /dev/null +++ b/tests/test_causalstress.py @@ -0,0 +1,138 @@ +from __future__ import annotations + +import subprocess +import sys +from pathlib import Path + +from dovla_cil.eval.causalstress import ( + CAUSALSTRESS_CATEGORIES, + CausalStressConfig, + compute_causalstress_metrics, + generate_causalstress_groups, +) +from dovla_cil.generation.pipeline import generate_cil_dataset +from dovla_cil.tasks.library import built_in_toy_tasks +from dovla_cil.training.trainer import DoVLATrainer, TrainerConfig +from dovla_cil.utils.io import read_json + + +def test_causalstress_generation_works() -> None: + groups = generate_causalstress_groups( + CausalStressConfig(num_tasks=len(CAUSALSTRESS_CATEGORIES), k=4, seed=1) + ) + assert len(groups) == len(CAUSALSTRESS_CATEGORIES) + assert {group.category for group in groups} == set(CAUSALSTRESS_CATEGORIES) + assert all(len(group.records) == 4 for group in groups) + assert all(record.group_id == group.group_id for group in groups for record in group.records) + + +def test_each_causalstress_category_generates_one_group() -> None: + for category in CAUSALSTRESS_CATEGORIES: + groups = generate_causalstress_groups( + CausalStressConfig(num_tasks=1, k=3, seed=4, categories=(category,)) + ) + assert len(groups) == 1 + assert groups[0].category == category + assert groups[0].records + assert groups[0].task.success_predicates + + +def test_hard_causalstress_categories_cycle_named_variants() -> None: + expected_counts = { + "similar_distractors": 4, + "spatial_relation_minimal_pairs": 3, + "negation_and_avoidance": 2, + "sequential_tasks": 3, + "irreversible_failure": 2, + "physics_perturbation_placeholders": 3, + } + for category, count in expected_counts.items(): + groups = generate_causalstress_groups( + CausalStressConfig(num_tasks=count, k=2, seed=5, categories=(category,)) + ) + assert len({group.task.task_id for group in groups}) == count + + +def test_causalstress_metrics_on_synthetic_predictions() -> None: + groups = generate_causalstress_groups(CausalStressConfig(num_tasks=3, k=4, seed=2)) + predictions = {} + for group in groups: + predictions[group.group_id] = { + "scores": [record.reward.score for record in group.records], + "success": [1.0 if record.reward.terminal_success else 0.0 for record in group.records], + "progress": [record.reward.progress for record in group.records], + "regret": [float(record.regret or 0.0) for record in group.records], + "effects": [ + [0.0] * 32 for _record in group.records + ], + } + # Use target effect vectors as predictions to make the MAE exact zero. + from dovla_cil.eval.causalstress import _effect_vector + + for group in groups: + predictions[group.group_id]["effects"] = [ + _effect_vector(record, dim=32) for record in group.records + ] + + metrics = compute_causalstress_metrics(groups, predictions) + assert metrics["pairwise_ranking_accuracy"] == 1.0 + assert metrics["top1_action_selection"] == 1.0 + assert metrics["success_prediction_accuracy"] == 1.0 + assert metrics["effect_prediction_mae"] == 0.0 + assert metrics["regret_calibration_error"] == 0.0 + assert "per_category" in metrics + assert "target_confusion_matrix" in metrics + + +def test_eval_causalstress_script_runs_on_smoke_checkpoint(tmp_path: Path) -> None: + dataset_dir = tmp_path / "cil" + run_dir = tmp_path / "run" + out_path = tmp_path / "causalstress.json" + generate_cil_dataset( + backend="toy", + tasks=built_in_toy_tasks()[:2], + out_dir=dataset_dir, + num_states_per_task=1, + k=4, + seed=3, + shard_size=8, + inline_observations=True, + ) + DoVLATrainer( + TrainerConfig( + dataset_dir=dataset_dir, + output_dir=run_dir, + epochs=1, + batch_groups=1, + records_per_group=4, + hidden_dim=32, + seed=3, + device="cpu", + ) + ).train() + + subprocess.run( + [ + sys.executable, + "scripts/eval_causalstress.py", + "--checkpoint", + str(run_dir / "best.pt"), + "--backend", + "toy", + "--out", + str(out_path), + "--num-tasks", + "3", + "--k", + "4", + "--seed", + "3", + ], + check=True, + capture_output=True, + text=True, + ) + metrics = read_json(out_path) + assert metrics["num_groups"] == 3 + assert "pairwise_ranking_accuracy" in metrics + assert "task_success_rate" in metrics diff --git a/tests/test_cil_images.py b/tests/test_cil_images.py new file mode 100644 index 0000000000000000000000000000000000000000..481a284cc60732bd5da614592e296e0578dfadef --- /dev/null +++ b/tests/test_cil_images.py @@ -0,0 +1,101 @@ +from __future__ import annotations + +import io +import json +from types import SimpleNamespace + +import numpy as np +import pytest + +from dovla_cil.data.images import CILImageReader +from dovla_cil.generation.pipeline import generate_cil_dataset +from dovla_cil.tasks.library import built_in_toy_tasks +from dovla_cil.training.trainer import DoVLATrainer, TrainerConfig + + +def test_image_reader_resolves_collection_source_and_caches_archive(tmp_path) -> None: + h5py = pytest.importorskip("h5py") + image_module = pytest.importorskip("PIL.Image") + source = tmp_path / "source" + source.mkdir() + expected = np.full((6, 8, 3), 117, dtype=np.uint8) + buffer = io.BytesIO() + image_module.fromarray(expected).save(buffer, format="JPEG", quality=95) + encoded = np.frombuffer(buffer.getvalue(), dtype=np.uint8) + with h5py.File(source / "observations.h5", "w") as handle: + dataset = handle.create_dataset( + "initial_rgb_jpeg", + shape=(1,), + dtype=h5py.vlen_dtype(np.dtype("uint8")), + ) + dataset[0] = encoded + record = SimpleNamespace( + record_id="record-0", + observation_ref="observations.h5#initial_rgb_jpeg/0", + next_observation_ref=None, + metadata={"source_dataset": str(source)}, + ) + + with CILImageReader(tmp_path / "collection") as reader: + actual = reader.read(record) + cached = reader.read(record) + assert len(reader._handles) == 1 + + assert actual.shape == expected.shape + assert actual.dtype == np.uint8 + assert np.array_equal(actual, cached) + + +def test_rgb_trainer_smoke_writes_checkpoint(tmp_path) -> None: + h5py = pytest.importorskip("h5py") + image_module = pytest.importorskip("PIL.Image") + torch = pytest.importorskip("torch") + dataset_dir = tmp_path / "data" + generate_cil_dataset( + backend="toy", + tasks=built_in_toy_tasks()[:2], + out_dir=dataset_dir, + num_states_per_task=1, + k=2, + seed=9, + shard_size=8, + inline_observations=True, + ) + image = np.zeros((24, 24, 3), dtype=np.uint8) + image[..., 0] = 180 + buffer = io.BytesIO() + image_module.fromarray(image).save(buffer, format="JPEG", quality=90) + encoded = np.frombuffer(buffer.getvalue(), dtype=np.uint8) + with h5py.File(dataset_dir / "observations.h5", "w") as handle: + images = handle.create_dataset( + "initial_rgb_jpeg", + shape=(1,), + dtype=h5py.vlen_dtype(np.dtype("uint8")), + ) + images[0] = encoded + for shard in (dataset_dir / "shards").glob("*.jsonl"): + records = [json.loads(line) for line in shard.read_text().splitlines()] + for record in records: + record["observation_ref"] = "observations.h5#initial_rgb_jpeg/0" + shard.write_text("".join(json.dumps(record) + "\n" for record in records)) + + run_dir = tmp_path / "run" + result = DoVLATrainer( + TrainerConfig( + dataset_dir=dataset_dir, + output_dir=run_dir, + epochs=1, + batch_groups=1, + records_per_group=2, + pair_count_per_group=1, + hidden_dim=24, + action_horizon=2, + effect_dim=8, + observation_mode="rgb", + device="cpu", + ) + ).train() + + checkpoint = torch.load(run_dir / "best.pt", map_location="cpu", weights_only=False) + assert result["best"] + assert checkpoint["model_config"]["observation_mode"] == "rgb" diff --git a/tests/test_cil_schema.py b/tests/test_cil_schema.py new file mode 100644 index 0000000000000000000000000000000000000000..1f9cc6f17a15b8578faa56adc17f2a5896f14df2 --- /dev/null +++ b/tests/test_cil_schema.py @@ -0,0 +1,98 @@ +from __future__ import annotations + +import pytest + +from dovla_cil.data.schema import ( + CIL_VERSION, + ActionChunk, + CILGroup, + CILRecord, + FailureInfo, + RewardInfo, + StructuredEffect, + compute_regret_and_ranks, + compute_state_hash, + make_record_id, + validate_group, +) + + +def make_record( + group_id: str, action_id: str, progress: float, *, state_hash: str = "s" +) -> CILRecord: + return CILRecord( + version=CIL_VERSION, + record_id=make_record_id(group_id, action_id, seed=7), + group_id=group_id, + state_hash=state_hash, + task_id="task", + scene_id=None, + instruction="move the mug", + instruction_family={"family": "place"}, + observation_ref=None, + observation_inline={"symbolic": True}, + action_chunk=ActionChunk( + action_id=action_id, + representation="delta_xy", + horizon=1, + values=[[progress, 0.0]], + skill_type="push", + ), + next_observation_ref=None, + next_observation_inline={"symbolic": True, "next": True}, + structured_effect=StructuredEffect( + object_pose_delta={"mug": [progress, 0.0, 0.0]}, + relation_before={"inside(mug,bowl)": False}, + relation_after={"inside(mug,bowl)": progress > 0.5}, + moved_objects=["mug"] if progress else [], + symbolic_before={"objects": {"mug": {"position": [0, 0, 0]}}}, + symbolic_after={"objects": {"mug": {"position": [progress, 0, 0]}}}, + ), + reward=RewardInfo( + progress=progress, + success=progress > 0.5, + terminal_success=progress > 0.5, + dense_components={"progress": progress}, + ), + regret=None, + rank_within_group=None, + candidate_type="sampled", + failure=None + if progress > 0 + else FailureInfo(type="no_motion", symbolic_reason="object did not move"), + ) + + +def test_schema_roundtrip() -> None: + record = make_record("g", "a", 1.0) + record.validate() + assert CILRecord.from_dict(record.to_dict()) == record + group = CILGroup.from_records([record]) + assert group.group_id == "g" + + +def test_group_validation() -> None: + records = [make_record("g", "a", 0.0), make_record("g", "b", 1.0)] + validate_group(records) + with pytest.raises(ValueError): + validate_group([make_record("g", "a", 0.0), make_record("other", "b", 1.0)]) + with pytest.raises(ValueError): + validate_group([make_record("g", "a", 0.0), make_record("g", "b", 1.0, state_hash="x")]) + + +def test_regret_and_rank_computation() -> None: + ranked = compute_regret_and_ranks( + [make_record("g", "low", 0.0), make_record("g", "high", 1.0)] + ) + by_action = {record.action_chunk.action_id: record for record in ranked} + assert by_action["high"].rank_within_group == 0 + assert by_action["high"].regret == 0.0 + assert by_action["low"].rank_within_group == 1 + assert by_action["low"].regret == 2.0 + + +def test_deterministic_record_ids_and_state_hashes() -> None: + assert make_record_id("g", "a", 1) == make_record_id("g", "a", 1) + assert make_record_id("g", "a", 1) != make_record_id("g", "a", 2) + assert compute_state_hash(b"state") == compute_state_hash(b"state") + assert compute_state_hash(b"state") != compute_state_hash(b"other") diff --git a/tests/test_config.py b/tests/test_config.py new file mode 100644 index 0000000000000000000000000000000000000000..53b05a7fe59b0c5521fe01fcb264557ff3d0a233 --- /dev/null +++ b/tests/test_config.py @@ -0,0 +1,117 @@ +from __future__ import annotations + +from pathlib import Path + +from dovla_cil.config.schema import ( + DoVLACILConfig, + apply_cli_overrides, + load_config, + load_config_dict, + save_resolved_config, +) + + +def test_load_default_config() -> None: + config = load_config() + assert isinstance(config, DoVLACILConfig) + assert config.sim.backend == "toy" + assert config.sim.seed == 0 + assert config.sim.params == {} + assert config.generation.backend == "toy" + assert config.generation.k > 0 + + +def test_env_overrides(monkeypatch) -> None: + monkeypatch.setenv("DOVLA_CIL_GENERATION__K", "12") + monkeypatch.setenv("DOVLA_CIL_SIM__SEED", "42") + monkeypatch.setenv("OPENCLAUDE_MODEL", "mock-model") + config = load_config() + assert config.generation.k == 12 + assert config.sim.seed == 42 + assert config.vlm.model == "mock-model" + + +def test_config_file_loads() -> None: + config = load_config("configs/toy/generate_cil_k4.yaml") + + assert config.generation.k == 4 + assert config.generation.output_dir == "outputs/config_toy/cil_k4" + assert config.sim.backend == "toy" + + +def test_all_experiment_configs_load_as_dicts() -> None: + paths = sorted(Path("configs").glob("*/*.yaml")) + assert paths + for path in paths: + payload = load_config_dict(path) + assert isinstance(payload, dict) + assert "seed" in payload + + +def test_large_templates_have_safe_defaults() -> None: + config = load_config("configs/large/generate_cil_maniskill_template.yaml") + + assert config.seed == 0 + assert config.sim.backend == "maniskill" + assert config.sim.params["env_id"] == "PickCube-v1" + assert config.generation.output_dir == "data/cil_maniskill_k32" + + +def test_env_expansion(monkeypatch, tmp_path: Path) -> None: + config_path = tmp_path / "config.yaml" + config_path.write_text( + "generation:\n" + " output_dir: ${DOVLA_TMP_ROOT}/cil\n" + " k: 4\n", + encoding="utf-8", + ) + monkeypatch.setenv("DOVLA_TMP_ROOT", str(tmp_path / "expanded")) + + payload = load_config_dict(config_path) + + assert payload["generation"]["output_dir"] == str(tmp_path / "expanded" / "cil") + + +def test_env_expansion_default_value(monkeypatch) -> None: + monkeypatch.delenv("DOVLA_MISSING_ROOT", raising=False) + payload = load_config_dict( + "configs/toy/generate_cil_k4.yaml", + overrides=["generation.output_dir=${DOVLA_MISSING_ROOT:-outputs/fallback}/cil"], + ) + + assert payload["generation"]["output_dir"] == "outputs/fallback/cil" + + +def test_cli_override_utility() -> None: + payload = load_config_dict( + "configs/toy/train_dovla_small.yaml", + overrides=["training.epochs=3", "model.hidden_dim=96", "sim.params.foo=bar", "sim.params.enabled=true"], # noqa: E501 + ) + config = DoVLACILConfig.from_dict(payload) + + assert config.training.epochs == 3 + assert config.model.hidden_dim == 96 + assert config.sim.params["foo"] == "bar" + assert config.sim.params["enabled"] is True + + +def test_typed_loader_accepts_cli_overrides() -> None: + config = load_config( + "configs/toy/train_dovla_small.yaml", + overrides=["training.epochs=2", "model.hidden_dim=80"], + ) + + assert config.training.epochs == 2 + assert config.model.hidden_dim == 80 + + +def test_apply_cli_overrides_and_save_resolved(tmp_path: Path) -> None: + payload = {"training": {"epochs": 1}, "paths": {"out": "${HOME}/runs"}} + apply_cli_overrides(payload, ["training.epochs=2", "paths.out=/tmp/run"]) + + path = save_resolved_config(payload, tmp_path) + + assert payload["training"]["epochs"] == 2 + assert payload["paths"]["out"] == "/tmp/run" + assert path == tmp_path / "resolved_config.yaml" + assert path.exists() diff --git a/tests/test_data_sharding.py b/tests/test_data_sharding.py new file mode 100644 index 0000000000000000000000000000000000000000..ee20ab8a5f20f5c9e8d01e95287d3d24218dd975 --- /dev/null +++ b/tests/test_data_sharding.py @@ -0,0 +1,185 @@ +from __future__ import annotations + +import subprocess +import sys +from pathlib import Path + +from dovla_cil.data.schema import ( + CIL_VERSION, + ActionChunk, + CILRecord, + RewardInfo, + StructuredEffect, + make_record_id, +) +from dovla_cil.data.sharding import ( + ShardReader, + ShardWriter, + iter_cil_records, + split_records_by_group, + write_cil_shards, +) +from dovla_cil.utils.io import iter_jsonl, read_json + + +def make_record(group_id: str, index: int, *, state_hash: str = "state-a") -> CILRecord: + action = ActionChunk( + action_id=f"action-{index}", + representation="delta_xy", + horizon=1, + values=[[float(index), 0.0]], + skill_type="push", + ) + return CILRecord( + version=CIL_VERSION, + record_id=make_record_id(group_id, action.action_id, seed=0), + group_id=group_id, + state_hash=state_hash, + task_id="task-a", + scene_id=None, + instruction="reach", + instruction_family={"family": "toy"}, + observation_ref=None, + observation_inline={"index": index}, + action_chunk=action, + next_observation_ref=None, + next_observation_inline={"index": index + 1}, + structured_effect=StructuredEffect( + object_pose_delta={"object": [1.0, 0.0, 0.0]}, + moved_objects=["object"], + ), + reward=RewardInfo( + progress=float(index), + success=index > 0, + terminal_success=index > 0, + dense_components={"progress": float(index)}, + ), + regret=None, + rank_within_group=None, + candidate_type="test", + failure=None, + ) + + +def test_split_keeps_groups_together() -> None: + records = [make_record("a", 0), make_record("b", 1), make_record("a", 2)] + shards = split_records_by_group(records, max_records_per_shard=2) + locations: dict[str, int] = {} + for shard_index, shard in enumerate(shards): + for record in shard: + if record.group_id in locations: + assert locations[record.group_id] == shard_index + locations[record.group_id] = shard_index + + +def test_write_shards_roundtrip(tmp_path: Path) -> None: + records = [make_record("a", 0), make_record("a", 1), make_record("b", 2)] + manifest = write_cil_shards(records, output_dir=tmp_path, max_records_per_shard=2) + restored = [] + for shard in manifest["shards"]: + restored.extend(iter_cil_records(tmp_path / str(shard["path"]))) + assert manifest["num_records"] == 3 + assert manifest["record_count"] == 3 + assert (tmp_path / "metadata.json").exists() + assert (tmp_path / "manifest.json").exists() + assert (tmp_path / "group_index.jsonl").exists() + assert (tmp_path / "record_index.jsonl").exists() + assert restored == records + + +def test_shard_writer_and_reader_roundtrip(tmp_path: Path) -> None: + records = [ + make_record("group-a", 0), + make_record("group-a", 1), + make_record("group-b", 2, state_hash="state-b"), + ] + writer = ShardWriter( + tmp_path, + dataset_name="toy_dataset", + backend="toy", + k=2, + task_count=1, + seed=7, + shard_size=2, + ) + for record in records: + writer.write(record) + metadata = writer.close() + + reader = ShardReader(tmp_path) + restored = list(reader.iterate_records()) + assert metadata["dataset_name"] == "toy_dataset" + assert metadata["backend"] == "toy" + assert metadata["num_records"] == 3 + assert metadata["num_groups"] == 2 + assert metadata["shards"][0]["path"] == "shards/shard_000000.jsonl" + assert restored == records + + +def test_group_index_is_correct_and_group_loads(tmp_path: Path) -> None: + records = [ + make_record("group-a", 0), + make_record("group-a", 1), + make_record("group-b", 2, state_hash="state-b"), + ] + write_cil_shards(records, output_dir=tmp_path, max_records_per_shard=2) + + group_index = {row["group_id"]: row for row in iter_jsonl(tmp_path / "group_index.jsonl")} + group_a = group_index["group-a"] + assert group_a["shard_path"] == "shards/shard_000000.jsonl" + assert group_a["record_ids"] == [records[0].record_id, records[1].record_id] + assert group_a["task_id"] == "task-a" + assert group_a["state_hash"] == "state-a" + assert group_a["num_records"] == 2 + assert group_a["max_reward"] == 1.0 + assert group_a["success_count"] == 1 + assert group_a["candidate_type_counts"] == {"test": 2} + + reader = ShardReader(tmp_path / "metadata.json") + assert reader.load_group("group-a") == records[:2] + assert list(reader.iterate_groups())[1] == [records[2]] + + +def test_metadata_layout_and_record_index(tmp_path: Path) -> None: + records = [make_record("group-a", 0), make_record("group-a", 1)] + write_cil_shards( + records, + output_dir=tmp_path, + max_records_per_shard=10, + dataset_name="cil_toy", + backend="toy", + k=2, + task_count=1, + seed=3, + ) + metadata = read_json(tmp_path / "metadata.json") + record_index = list(iter_jsonl(tmp_path / "record_index.jsonl")) + + assert metadata["dataset_name"] == "cil_toy" + assert metadata["version"] + assert metadata["created_at"] + assert metadata["backend"] == "toy" + assert metadata["num_groups"] == 1 + assert metadata["num_records"] == 2 + assert metadata["k"] == 2 + assert metadata["task_count"] == 1 + assert metadata["seed"] == 3 + assert metadata["schema_version"] == CIL_VERSION + assert record_index[0]["record_id"] == records[0].record_id + assert record_index[0]["shard_path"] == "shards/shard_000000.jsonl" + + +def test_inspect_script_prints_summary_and_ranking(tmp_path: Path) -> None: + records = [make_record("group-a", 0), make_record("group-a", 1)] + write_cil_shards(records, output_dir=tmp_path, max_records_per_shard=10) + + result = subprocess.run( + [sys.executable, "scripts/inspect_shard.py", str(tmp_path), "--group-id", "group-a"], + check=True, + text=True, + capture_output=True, + ) + + assert "num_records: 2" in result.stdout + assert "sample_group: group-a" in result.stdout + assert "record_id\tcandidate_type\treward.progress\tsuccess\tregret\trank\tfailure.type" in result.stdout diff --git a/tests/test_dataset_reports.py b/tests/test_dataset_reports.py new file mode 100644 index 0000000000000000000000000000000000000000..739ba734667cebe39c37a6ea362d9bf5b174c67c --- /dev/null +++ b/tests/test_dataset_reports.py @@ -0,0 +1,182 @@ +from __future__ import annotations + +import csv +import subprocess +import sys +from pathlib import Path + +from dovla_cil.data.schema import ( + CIL_VERSION, + ActionChunk, + CILRecord, + FailureInfo, + RewardInfo, + StructuredEffect, + make_record_id, +) +from dovla_cil.data.sharding import write_cil_shards +from dovla_cil.experiments.reports import ( + compute_candidate_stats, + compute_failure_stats, + generate_dataset_report, + load_dataset_summary, +) +from dovla_cil.utils.io import read_json + + +def test_dataset_report_script_runs_and_creates_files(tmp_path: Path) -> None: + dataset_dir = tmp_path / "dataset" + out_dir = tmp_path / "report" + _write_tiny_dataset(dataset_dir) + + result = subprocess.run( + [ + sys.executable, + "scripts/report_dataset.py", + "--dataset", + str(dataset_dir), + "--out", + str(out_dir), + "--sample-groups", + "2", + ], + check=True, + text=True, + capture_output=True, + ) + + expected = { + "summary.json", + "candidate_type_counts.csv", + "reward_histogram.png", + "success_by_candidate_type.csv", + "success_by_candidate_type.png", + "regret_histogram.png", + "group_size_distribution.png", + "failure_type_counts.csv", + "failure_type_counts.png", + "examples.md", + } + assert "num records: 3" in result.stdout + assert expected.issubset({path.name for path in out_dir.iterdir()}) + + +def test_report_stats_match_expected_toy_dataset(tmp_path: Path) -> None: + dataset_dir = tmp_path / "dataset" + out_dir = tmp_path / "report" + records = _write_tiny_dataset(dataset_dir) + + summary = generate_dataset_report(dataset_dir, out_dir, sample_groups=2, seed=4) + candidate_stats = { + row["candidate_type"]: row for row in compute_candidate_stats(records) + } + failure_stats = {row["failure_type"]: row for row in compute_failure_stats(records)} + summary_file = read_json(out_dir / "summary.json") + candidate_counts = _read_csv(out_dir / "candidate_type_counts.csv") + + assert summary["num_records"] == 3 + assert summary["num_groups"] == 2 + assert summary["candidate_type_counts"] == {"expert": 2, "wrong_target": 1} + assert summary["failure_type_counts"] == {"success": 2, "wrong_target": 1} + assert summary_file["success_rate"] == 2 / 3 + assert candidate_stats["expert"]["success_rate"] == 1.0 + assert candidate_stats["wrong_target"]["success_rate"] == 0.0 + assert failure_stats["wrong_target"]["count"] == 1 + assert candidate_counts == [ + {"candidate_type": "expert", "count": "2"}, + {"candidate_type": "wrong_target", "count": "1"}, + ] + assert "| record_id | candidate_type | reward.progress |" in ( + out_dir / "examples.md" + ).read_text(encoding="utf-8") + + +def test_load_dataset_summary_function(tmp_path: Path) -> None: + dataset_dir = tmp_path / "dataset" + _write_tiny_dataset(dataset_dir) + + summary = load_dataset_summary(dataset_dir) + + assert summary["num_records"] == 3 + assert summary["reward"]["min"] == 0.2 + assert summary["reward"]["max"] == 1.0 + assert summary["group_size"]["max"] == 2 + + +def _write_tiny_dataset(dataset_dir: Path) -> list[CILRecord]: + records = [ + _record("group-a", 0, "expert", 1.0, True, "success", regret=0.0, rank=0), + _record("group-a", 1, "wrong_target", 0.2, False, "wrong_target", regret=0.8, rank=1), + _record("group-b", 0, "expert", 1.0, True, "success", regret=0.0, rank=0), + ] + write_cil_shards( + records, + output_dir=dataset_dir, + max_records_per_shard=10, + dataset_name="report_toy", + backend="toy", + k=2, + task_count=1, + seed=5, + ) + return records + + +def _record( + group_id: str, + index: int, + candidate_type: str, + reward_progress: float, + success: bool, + failure_type: str, + *, + regret: float, + rank: int, +) -> CILRecord: + action = ActionChunk( + action_id=f"{candidate_type}-{index}", + representation="semantic", + horizon=1, + values=[{"command": "noop" if not success else "grasp", "object": "red_mug"}], + skill_type="grasp", + metadata={"candidate_type": candidate_type, "intended_target": "red_mug"}, + ) + return CILRecord( + version=CIL_VERSION, + record_id=make_record_id(group_id, action.action_id, seed=9), + group_id=group_id, + state_hash=f"state-{group_id}", + task_id="toy_report_task", + scene_id=f"scene-{group_id}", + instruction="Pick the red mug.", + instruction_family={"family": "pick"}, + observation_ref=None, + observation_inline={"symbolic_state": {"objects": {}}}, + action_chunk=action, + next_observation_ref=None, + next_observation_inline={"symbolic_state": {"objects": {}}}, + structured_effect=StructuredEffect( + object_pose_delta={"red_mug": [reward_progress, 0.0, 0.0]}, + relation_after={"grasped(red_mug)": success}, + moved_objects=["red_mug"] if reward_progress > 0 else [], + ), + reward=RewardInfo( + progress=reward_progress, + success=success, + terminal_success=success, + dense_components={"progress": reward_progress}, + ), + regret=regret, + rank_within_group=rank, + candidate_type=candidate_type, + failure=FailureInfo( + type=failure_type, + symbolic_reason=failure_type, + language_explanation=failure_type, + ), + ) + + +def _read_csv(path: Path) -> list[dict[str, str]]: + with path.open("r", encoding="utf-8", newline="") as handle: + return list(csv.DictReader(handle)) diff --git a/tests/test_distributed_generation.py b/tests/test_distributed_generation.py new file mode 100644 index 0000000000000000000000000000000000000000..be789c6aa85388c30b4a63177eed3b7c00256fe5 --- /dev/null +++ b/tests/test_distributed_generation.py @@ -0,0 +1,70 @@ +from __future__ import annotations + +import importlib.util +from pathlib import Path + +import pytest + +from dovla_cil.generation.distributed import ( + DistributedCILConfig, + plan_generation_jobs, + require_ray, + run_distributed_cil_generation, +) +from dovla_cil.tasks.library import built_in_toy_tasks + + +def test_ray_absent_gives_clear_error() -> None: + if importlib.util.find_spec("ray") is not None: + pytest.skip("Ray is installed in this environment") + + with pytest.raises(ImportError) as exc_info: + require_ray() + + assert "Ray is optional" in str(exc_info.value) + assert "pip install ray" in str(exc_info.value) + + +def test_deterministic_group_ids_under_seed() -> None: + tasks = built_in_toy_tasks()[:2] + + first = plan_generation_jobs(tasks, backend="toy", num_states_per_task=2, seed=7) + second = plan_generation_jobs(tasks, backend="toy", num_states_per_task=2, seed=7) + skipped = plan_generation_jobs( + tasks, + backend="toy", + num_states_per_task=2, + seed=7, + completed_group_ids={first[0].group_id}, + ) + + assert [job.group_id for job in first] == [job.group_id for job in second] + assert len(first) == 4 + assert len(skipped) == 3 + assert first[0].group_id not in {job.group_id for job in skipped} + + +def test_tiny_distributed_generation_if_ray_present(tmp_path: Path) -> None: + ray = pytest.importorskip("ray") + tasks = built_in_toy_tasks()[:1] + config = DistributedCILConfig( + backend="toy", + output_dir=tmp_path, + num_workers=1, + num_states_per_task=1, + k=2, + seed=3, + shard_size=10, + ) + + try: + summary = run_distributed_cil_generation(config, tasks) + finally: + if ray.is_initialized(): + ray.shutdown() + + assert summary.num_groups == 1 + assert summary.num_records == 2 + assert (tmp_path / "manifest.json").exists() + assert (tmp_path / "distributed_manifest.json").exists() + assert (tmp_path / "group_index.jsonl").exists() diff --git a/tests/test_dovla_model.py b/tests/test_dovla_model.py new file mode 100644 index 0000000000000000000000000000000000000000..a0b1effe2e97ecf3e153077dc59962923578b21a --- /dev/null +++ b/tests/test_dovla_model.py @@ -0,0 +1,163 @@ +from __future__ import annotations + +import pytest + +torch = pytest.importorskip("torch") + +from dovla_cil.data.schema import ActionChunk +from dovla_cil.models import devectorize_toy_action, vectorize_toy_action +from dovla_cil.models.dovla import DoVLAConfig, DoVLAModel, load_model_state + + +def _model_and_inputs(): + config = DoVLAConfig( + obs_dim=10, + lang_dim=16, + action_dim=8, + hidden_dim=32, + action_horizon=3, + effect_dim=7, + intervention_dim=24, + ) + model = DoVLAModel(config) + observation = torch.randn(4, config.obs_dim) + instructions = [ + "pick the red mug", + "put the mug in the bowl", + "open the drawer", + "push the cube", + ] + action = torch.randn(4, config.action_horizon, config.action_dim) + return model, config, observation, instructions, action + + +def test_forward_policy_works_and_shapes() -> None: + model, config, observation, instructions, _action = _model_and_inputs() + predicted = model.forward_policy(observation, instructions) + assert predicted.shape == (4, config.action_horizon, config.action_dim) + + +def test_forward_effect_works_and_shapes() -> None: + model, config, observation, instructions, action = _model_and_inputs() + output = model.forward_effect(observation, instructions, action) + assert output["effect_vector"].shape == (4, config.effect_dim) + assert output["success_logit"].shape == (4,) + assert output["success"].shape == (4,) + assert output["progress"].shape == (4,) + + +def test_forward_reward_works_and_shapes() -> None: + model, _config, observation, instructions, action = _model_and_inputs() + reward = model.forward_reward(observation, instructions, action) + regret = model.forward_regret(observation, instructions, action) + z = model.encode_intervention(observation, instructions, action) + assert reward.shape == (4,) + assert regret.shape == (4,) + assert z.shape == (4, model.config.intervention_dim) + + +def test_gradients_flow() -> None: + model, _config, observation, instructions, action = _model_and_inputs() + policy = model.forward_policy(observation, instructions) + effect = model.forward_effect(observation, instructions, action) + reward = model.forward_reward(observation, instructions, action) + loss = policy.square().mean() + effect["effect_vector"].square().mean() + reward.mean() + loss.backward() + grads = [ + parameter.grad + for parameter in model.parameters() + if parameter.requires_grad and parameter.grad is not None + ] + assert grads + assert sum(float(grad.abs().sum()) for grad in grads) > 0.0 + + +def test_rgb_observation_encoder_drives_policy_and_field_gradients() -> None: + config = DoVLAConfig( + obs_dim=10, + lang_dim=16, + action_dim=7, + hidden_dim=32, + action_horizon=2, + effect_dim=6, + intervention_dim=24, + observation_mode="rgb", + ) + model = DoVLAModel(config) + images = torch.randint(0, 256, (3, 32, 40, 3), dtype=torch.uint8) + instructions = ["push the cube", "stack the cubes", "lift the peg"] + actions = torch.randn(3, config.action_horizon, config.action_dim) + + policy = model.forward_policy(images, instructions) + field = model.forward_field(images, instructions, actions) + (policy.square().mean() + field["potential"].mean()).backward() + + assert policy.shape == (3, 2, 7) + assert field["effect_vector"].shape == (3, 6) + assert any( + parameter.grad is not None and float(parameter.grad.abs().sum()) > 0 + for parameter in model.observation_encoder.image_net.parameters() + ) + + +def test_toy_action_vectorize_and_devectorize() -> None: + action = ActionChunk( + representation="semantic", + horizon=2, + values=[ + {"command": "move_to", "object": "red_mug"}, + {"command": "push", "object": "red_mug", "dx": 0.1, "dy": -0.2}, + ], + skill_type="push", + ) + matrix = vectorize_toy_action(action, action_dim=8, action_horizon=3) + restored = devectorize_toy_action(matrix, skill_type="push") + assert len(matrix) == 3 + assert len(matrix[0]) == 8 + assert restored.representation == "semantic" + assert restored.skill_type == "push" + + +def test_numeric_action_chunk_vectorization_preserves_simulator_controls() -> None: + action = ActionChunk( + representation="maniskill_pd_ee_delta_pose", + horizon=2, + values=[ + [0.1, -0.2, 0.3, 0.4, -0.5, 0.6, 0.7], + [0.0, 0.1, -0.1, 0.2, -0.2, 0.3, -0.3], + ], + skill_type="pick_cube", + ) + + matrix = vectorize_toy_action(action, action_dim=8, action_horizon=3) + + assert matrix[0] == [0.1, -0.2, 0.3, 0.4, -0.5, 0.6, 0.7, 0.0] + assert matrix[1] == [0.0, 0.1, -0.1, 0.2, -0.2, 0.3, -0.3, 0.0] + assert matrix[2] == [0.0] * 8 + + +def test_model_state_loader_allows_only_declared_omitted_prefixes() -> None: + model, config, _observation, _instructions, _action = _model_and_inputs() + state = model.state_dict() + prefix = "observation_encoder.net.0." + compact = {key: value for key, value in state.items() if not key.startswith(prefix)} + restored = DoVLAModel(config) + + load_model_state( + restored, + { + "model_state_dict": compact, + "omitted_state_prefixes": [prefix], + }, + ) + + invalid = dict(compact) + invalid.pop("language_encoder.net.0.weight") + with pytest.raises(RuntimeError, match="invalid_missing"): + load_model_state( + DoVLAModel(config), + { + "model_state_dict": invalid, + "omitted_state_prefixes": [prefix], + }, + ) diff --git a/tests/test_effect_rewards.py b/tests/test_effect_rewards.py new file mode 100644 index 0000000000000000000000000000000000000000..788c8c352a6c4bf74b41abeeed8a3d25d117a4a6 --- /dev/null +++ b/tests/test_effect_rewards.py @@ -0,0 +1,152 @@ +from __future__ import annotations + +import pytest + +from dovla_cil.data.schema import ActionChunk +from dovla_cil.effects.extractors import extract_structured_effect +from dovla_cil.effects.failure_classifier import classify_failure, classify_toy_failure +from dovla_cil.effects.rewards import best_action_index, compute_reward, distance_to_target_reward +from dovla_cil.tasks.library import ToyTaskLibrary + + +def test_distance_reward_prefers_near_actions() -> None: + assert distance_to_target_reward(0.1, success=False) > distance_to_target_reward( + 0.5, success=False + ) + assert distance_to_target_reward(0.01, success=True) > 0.0 + + +def test_distance_reward_rejects_negative() -> None: + with pytest.raises(ValueError): + distance_to_target_reward(-1.0, success=False) + + +def test_success_predicate_gives_success_reward() -> None: + task = ToyTaskLibrary().get_by_id("toy_put_red_mug_in_blue_bowl") + before = _state(red_mug=[0.0, 0.0, 0.03], blue_bowl=[1.0, 0.0, 0.03]) + after = _state( + red_mug=[1.0, 0.0, 0.04], + blue_bowl=[1.0, 0.0, 0.03], + red_mug_extra={"inside": "blue_bowl"}, + ) + + effect = extract_structured_effect(before, after, task=task) + reward = compute_reward(task, effect) + + assert reward.success is True + assert reward.terminal_success is True + assert reward.progress == 1.0 + + +def test_wrong_target_classified() -> None: + task = ToyTaskLibrary().get_by_id("toy_pick_object_among_distractors") + before = _state(red_mug=[0, 0, 0.03], blue_mug=[0.5, 0, 0.03], green_bowl=[1, 0, 0.03]) + after = _state(red_mug=[0, 0, 0.03], blue_mug=[0.8, 0, 0.03], green_bowl=[1, 0, 0.03]) + effect = extract_structured_effect(before, after, task=task) + reward = compute_reward(task, effect) + action = ActionChunk( + representation="semantic", + values=[{"command": "grasp", "object": "blue_mug"}], + skill_type="grasp", + metadata={"candidate_type": "wrong_target", "intended_target": "blue_mug"}, + ) + + failure = classify_failure(task, action, effect, reward) + + assert failure.type == "wrong_target" + + +def test_no_motion_classified() -> None: + task = ToyTaskLibrary().get_by_id("toy_lift_can") + before = _state(can=[0, 0, 0.03]) + after = _state(can=[0, 0, 0.03]) + effect = extract_structured_effect(before, after, task=task) + reward = compute_reward(task, effect) + action = ActionChunk( + representation="semantic", + values=[{"command": "noop"}], + skill_type="noop", + metadata={"candidate_type": "noop", "intended_target": "can"}, + ) + + failure = classify_failure(task, action, effect, reward) + + assert failure.type == "no_motion" + + +def test_partial_pick_place_classified() -> None: + task = ToyTaskLibrary().get_by_id("toy_put_red_mug_in_blue_bowl") + before = _state(red_mug=[0.0, 0.0, 0.03], blue_bowl=[1.0, 0.0, 0.03]) + after = _state( + red_mug=[0.5, 0.0, 0.2], + blue_bowl=[1.0, 0.0, 0.03], + red_mug_extra={"grasped": True, "lifted": True}, + ) + effect = extract_structured_effect( + before, after, task=task, rollout_info={"grasp_success": True} + ) + reward = compute_reward(task, effect) + action = ActionChunk( + representation="semantic", + values=[{"command": "grasp", "object": "red_mug"}], + skill_type="grasp", + metadata={"candidate_type": "near_miss", "intended_target": "red_mug"}, + ) + + failure = classify_failure(task, action, effect, reward) + + assert reward.success is False + assert 0.0 < reward.progress < 1.0 + assert failure.type == "partial_success" + + +def test_progress_is_bounded() -> None: + task = ToyTaskLibrary().get_by_id("toy_push_cube_to_target_zone") + before = _state(cube=[0, 0, 0.03], target_zone=[10, 0, 0.0]) + after = _state(cube=[100, 0, 0.03], target_zone=[10, 0, 0.0]) + effect = extract_structured_effect(before, after, task=task) + reward = compute_reward(task, effect) + + assert 0.0 <= reward.progress <= 1.0 + + +def test_effect_extraction_compatibility_and_toy_failure() -> None: + effect = extract_structured_effect( + {"position": 0.0}, {"position": 0.5}, info={"success": False} + ) + assert effect.metrics["delta_position"] == 0.5 + assert effect.relation_after["success"] is False + assert classify_toy_failure(distance=0.2, tolerance=0.05) == "missed_target" + assert classify_toy_failure(distance=0.01, tolerance=0.05) is None + + +def test_best_action_index() -> None: + assert best_action_index([0.0, 2.0, 1.0]) == 1 + + +def _state( + *, + red_mug=None, + blue_bowl=None, + blue_mug=None, + green_bowl=None, + can=None, + cube=None, + target_zone=None, + red_mug_extra=None, +) -> dict: + objects = {} + for name, position in { + "red_mug": red_mug, + "blue_bowl": blue_bowl, + "blue_mug": blue_mug, + "green_bowl": green_bowl, + "can": can, + "cube": cube, + "target_zone": target_zone, + }.items(): + if position is not None: + objects[name] = {"position": list(position), "grasped": False, "lifted": False} + if red_mug_extra and "red_mug" in objects: + objects["red_mug"].update(red_mug_extra) + return {"objects": objects, "near_threshold": 0.25, "lifted_z": 0.1} diff --git a/tests/test_eval_reports.py b/tests/test_eval_reports.py new file mode 100644 index 0000000000000000000000000000000000000000..d327f3cfd42211b920911b0748f5057e40eeafc7 --- /dev/null +++ b/tests/test_eval_reports.py @@ -0,0 +1,107 @@ +from __future__ import annotations + +import csv +import subprocess +import sys +from pathlib import Path + +from dovla_cil.experiments.reports import generate_eval_report +from dovla_cil.utils.io import write_json + + +def test_report_eval_works_with_fake_metrics(tmp_path: Path) -> None: + metrics_dir = tmp_path / "runs" + _write_fake_metrics(metrics_dir) + out_dir = tmp_path / "report" + + result = subprocess.run( + [ + sys.executable, + "scripts/report_eval.py", + "--inputs", + str(metrics_dir / "*" / "metrics.json"), + "--out", + str(out_dir), + "--name", + "fake_scaling", + ], + check=True, + text=True, + capture_output=True, + ) + + assert "num runs: 3" in result.stdout + assert (out_dir / "aggregate_metrics.csv").exists() + assert (out_dir / "report.md").exists() + assert (out_dir / "success_rate.png").exists() + assert (out_dir / "ranking_accuracy.png").exists() + assert (out_dir / "score_vs_k.png").exists() + + +def test_eval_report_markdown_and_csv_content(tmp_path: Path) -> None: + metrics_dir = tmp_path / "runs" + _write_fake_metrics(metrics_dir) + out_dir = tmp_path / "report" + + summary = generate_eval_report( + [metrics_dir / "*" / "metrics.json"], + out_dir, + experiment_name="fake_scaling", + ) + rows = _read_csv(out_dir / "aggregate_metrics.csv") + markdown = (out_dir / "report.md").read_text(encoding="utf-8") + + assert summary["num_runs"] == 3 + assert rows[0]["k"] == "1" + assert rows[1]["k"] == "2" + assert rows[2]["k"] == "4" + assert rows[1]["ranking_acc"] == "0.7" + assert "# fake_scaling" in markdown + assert "Best K by ranking_acc: `2`" in markdown + assert "Best K by success: `4`" in markdown + assert "beta_log_k" in markdown + assert "Warning: ranking_acc does not improve monotonically with K." in markdown + + +def _write_fake_metrics(root: Path) -> None: + payloads = [ + { + "run_name": "k1", + "k": 1, + "task_success_rate": 0.30, + "pairwise_ranking_accuracy": 0.50, + "top1_action_selection": 0.40, + "instruction_switch_accuracy": 0.20, + "effect_prediction_mae": 0.90, + "regret_calibration_error": 0.30, + }, + { + "run_name": "k2", + "k": 2, + "task_success_rate": 0.50, + "pairwise_ranking_accuracy": 0.70, + "top1_action_selection": 0.60, + "instruction_switch_accuracy": 0.40, + "effect_prediction_mae": 0.70, + "regret_calibration_error": 0.20, + }, + { + "run_name": "k4", + "k": 4, + "success_rate": 0.60, + "ranking_acc": 0.65, + "top1_action_selection": 0.70, + "instruction_switch_acc": 0.50, + "effect_mae": 0.60, + "regret_ece": 0.15, + "regression": {"ranking_acc": {"beta_log_k": 0.1}}, + }, + ] + for payload in payloads: + out = root / f"k_{int(payload['k']):04d}" / "metrics.json" + write_json(payload, out) + + +def _read_csv(path: Path) -> list[dict[str, str]]: + with path.open("r", encoding="utf-8", newline="") as handle: + return list(csv.DictReader(handle)) diff --git a/tests/test_external_checkpoint_verify.py b/tests/test_external_checkpoint_verify.py new file mode 100644 index 0000000000000000000000000000000000000000..1aa6c87ddd365f68263d4b7e81250bf684086733 --- /dev/null +++ b/tests/test_external_checkpoint_verify.py @@ -0,0 +1,59 @@ +from __future__ import annotations + +import subprocess +import sys +from pathlib import Path + +from scripts.verify_external_checkpoint import build_manifest + + +def test_external_checkpoint_manifest_ready_when_required_files_exist(tmp_path: Path) -> None: + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + (checkpoint / "config.json").write_text("{}", encoding="utf-8") + (checkpoint / "model.safetensors").write_bytes(b"weights") + + manifest = build_manifest( + checkpoint, + model_family="smolvla", + repo_id="lerobot/smolvla_base", + revision="abc", + required_files=["config.json", "model.safetensors"], + ) + + assert manifest["ready"] is True + assert manifest["missing_required_files"] == [] + assert manifest["file_count"] == 2 + assert manifest["total_bytes"] == len("{}") + len(b"weights") + model_row = next(row for row in manifest["files"] if row["path"] == "model.safetensors") + assert len(model_row["sha256"]) == 64 + + +def test_external_checkpoint_cli_reports_missing_required_files(tmp_path: Path) -> None: + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + (checkpoint / "config.json").write_text("{}", encoding="utf-8") + out = tmp_path / "manifest.json" + + result = subprocess.run( + [ + sys.executable, + "scripts/verify_external_checkpoint.py", + "--checkpoint", + str(checkpoint), + "--out", + str(out), + "--model-family", + "smolvla", + "--required-file", + "config.json", + "--required-file", + "model.safetensors", + ], + capture_output=True, + text=True, + ) + + assert result.returncode == 2 + assert "model.safetensors" in result.stdout + assert out.exists() diff --git a/tests/test_external_vla_baseline.py b/tests/test_external_vla_baseline.py new file mode 100644 index 0000000000000000000000000000000000000000..5a299a841d0920f04e95769fd2e651c7c4c9a482 --- /dev/null +++ b/tests/test_external_vla_baseline.py @@ -0,0 +1,134 @@ +from __future__ import annotations + +import json +import subprocess +import sys +from pathlib import Path + +from dovla_cil.eval.external_vla_baseline import ( + ExternalVLABaselineSpec, + assess_external_vla_baseline, + build_external_vla_plan, + redact_command, + run_external_vla_entrypoint, + write_external_vla_plan, +) + + +def test_external_vla_plan_is_secret_free_and_not_ready_without_adapter(tmp_path: Path) -> None: + dataset = tmp_path / "dataset" + dataset.mkdir() + spec = ExternalVLABaselineSpec( + model_family="smolvla", + checkpoint_path=str(tmp_path / "missing-checkpoint"), + dataset_dir=str(dataset), + out_dir=str(tmp_path / "run"), + ) + + status = assess_external_vla_baseline(spec) + plan = build_external_vla_plan(spec) + payload = json.dumps(plan) + + assert not status.ready + assert "adapter_entrypoint" in status.missing + assert "hf download lerobot/smolvla_base" in plan["commands"]["download"] + assert "TOKEN" not in payload + assert "OPENCLAUDE_API_KEY" not in payload + + +def test_external_vla_redacts_token_like_command_fragments() -> None: + command = "HF_TOKEN=abc hf download repo --token xyz --api-key nope" + + redacted = redact_command(command) + + assert "abc" not in redacted + assert "xyz" not in redacted + assert "nope" not in redacted + assert "" in redacted + + +def test_external_vla_dry_run_cli_writes_plan(tmp_path: Path) -> None: + dataset = tmp_path / "dataset" + dataset.mkdir() + out = tmp_path / "out" + + result = subprocess.run( + [ + sys.executable, + "scripts/run_external_vla_baseline.py", + "--model-family", + "smolvla", + "--checkpoint", + str(tmp_path / "missing-checkpoint"), + "--dataset", + str(dataset), + "--out", + str(out), + "--dry-run", + ], + check=True, + capture_output=True, + text=True, + ) + + assert "external VLA plan" in result.stdout + assert (out / "external_vla_baseline_plan.json").exists() + + +def test_external_vla_entrypoint_contract(tmp_path: Path, monkeypatch) -> None: + module_path = tmp_path / "fake_external_adapter.py" + module_path.write_text( + "def run(spec, plan):\n" + " return {'model_family': spec['model_family'], 'success_rate': 0.25}\n", + encoding="utf-8", + ) + monkeypatch.syspath_prepend(str(tmp_path)) + dataset = tmp_path / "dataset" + checkpoint = tmp_path / "checkpoint" + dataset.mkdir() + checkpoint.mkdir() + spec = ExternalVLABaselineSpec( + model_family="smolvla", + checkpoint_path=str(checkpoint), + dataset_dir=str(dataset), + out_dir=str(tmp_path / "run"), + adapter_entrypoint="fake_external_adapter:run", + ) + + plan_path = write_external_vla_plan(spec) + result = run_external_vla_entrypoint("fake_external_adapter:run", spec) + + assert plan_path.exists() + assert result == {"model_family": "smolvla", "success_rate": 0.25} + + +def test_external_vla_cli_passes_secret_free_adapter_config(tmp_path: Path) -> None: + dataset = tmp_path / "dataset" + checkpoint = tmp_path / "checkpoint" + out = tmp_path / "out" + dataset.mkdir() + checkpoint.mkdir() + adapter_config = tmp_path / "adapter.json" + adapter_config.write_text('{"steps": 2, "state_dim": 32}', encoding="utf-8") + + subprocess.run( + [ + sys.executable, + "scripts/run_external_vla_baseline.py", + "--checkpoint", + str(checkpoint), + "--dataset", + str(dataset), + "--out", + str(out), + "--adapter-config", + str(adapter_config), + "--dry-run", + ], + check=True, + capture_output=True, + text=True, + ) + + plan = json.loads((out / "external_vla_baseline_plan.json").read_text(encoding="utf-8")) + assert plan["spec"]["metadata"] == {"steps": 2, "state_dim": 32} diff --git a/tests/test_generation_pipeline.py b/tests/test_generation_pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..90179d811c3b1a8fe99b2dccb82ba132c2cd9222 --- /dev/null +++ b/tests/test_generation_pipeline.py @@ -0,0 +1,52 @@ +from __future__ import annotations + +from pathlib import Path + +from dovla_cil.data.sharding import iter_cil_records +from dovla_cil.generation.pipeline import generate_builtin_toy_dataset, generate_cil_dataset +from dovla_cil.tasks.library import built_in_toy_tasks +from dovla_cil.utils.io import iter_jsonl, read_json + + +def test_toy_generation_pipeline_writes_shards_and_group_index(tmp_path: Path) -> None: + summary = generate_cil_dataset( + backend="toy", + tasks=built_in_toy_tasks()[:2], + out_dir=tmp_path, + num_states_per_task=1, + k=4, + seed=11, + shard_size=4, + inline_observations=True, + ) + + manifest = read_json(tmp_path / "manifest.json") + group_index = list(iter_jsonl(tmp_path / "group_index.jsonl")) + records = [] + for shard in manifest["shards"]: + records.extend(iter_cil_records(tmp_path / str(shard["path"]))) + + assert summary.num_groups == 2 + assert summary.num_records == 8 + assert manifest["record_count"] == 8 + assert manifest["group_count"] == 2 + assert manifest["group_index_path"] == "group_index.jsonl" + assert len(group_index) == 2 + assert all(entry["num_records"] == 4 for entry in group_index) + assert all(record.rank_within_group is not None for record in records) + assert all(record.regret is not None for record in records) + assert any(record.reward.terminal_success for record in records) + assert "expert" in summary.candidate_type_distribution + + +def test_builtin_group_shortcut_does_not_cap_at_library_size(tmp_path: Path) -> None: + summary = generate_builtin_toy_dataset( + out_dir=tmp_path, + groups=12, + k=2, + seed=17, + shard_size=64, + ) + + assert summary.num_groups == 12 + assert summary.num_records == 24 diff --git a/tests/test_hpc_clean_results_report.py b/tests/test_hpc_clean_results_report.py new file mode 100644 index 0000000000000000000000000000000000000000..d57de0d7006826fc6c9795b6ea30ec7f635373d6 --- /dev/null +++ b/tests/test_hpc_clean_results_report.py @@ -0,0 +1,312 @@ +from __future__ import annotations + +import csv +import json +import subprocess +import sys +from pathlib import Path + + +def _write_eval(path: Path, **overrides: object) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + payload = { + "objective": "lattice_field", + "observation_mode": "state", + "training_k": 16, + "evaluation_k": 16, + "seed": 0, + "num_groups": 10, + "num_records": 160, + "num_pairs": 1000, + "pairwise_ranking_accuracy": 0.7, + "top1_action_selection": 0.4, + "selected_success_rate": 0.3, + "oracle_success_rate": 0.5, + "ndcg_at_k": 0.9, + "effect_prediction_mae": 0.03, + "selection_regret": 0.2, + "potential_edge_mae": 0.3, + } + payload.update(overrides) + path.write_text(json.dumps(payload), encoding="utf-8") + + +def test_clean_hpc_report_excludes_unclean_paths(tmp_path: Path) -> None: + root = tmp_path / "experiments" + _write_eval( + root + / "maniskill_presuccess_scaling_fixed14k" + / "runs" + / "k_1" + / "seed_0" + / "lattice_eval.json", + training_k=1, + pairwise_ranking_accuracy=0.55, + ) + _write_eval( + root + / "maniskill_presuccess_scaling_fixed14k" + / "runs" + / "k_16" + / "seed_0" + / "lattice_eval.json", + training_k=16, + pairwise_ranking_accuracy=0.70, + ) + _write_eval( + root + / "maniskill_presuccess_six_task_runs" + / "legacy" + / "seed_0" + / "lattice_eval.json", + objective="legacy", + pairwise_ranking_accuracy=0.64, + selected_success_rate=0.22, + ) + _write_eval( + root + / "maniskill_presuccess_six_task_runs" + / "lattice_field" + / "seed_0" + / "lattice_eval.json", + objective="lattice_field", + pairwise_ranking_accuracy=0.69, + selected_success_rate=0.28, + ) + _write_eval( + root / "presuccess_pick_pilot_runs" / "lattice_field" / "lattice_eval.json", + pairwise_ranking_accuracy=0.99, + ) + _write_eval( + root + / "maniskill_presuccess_transfer_leave_stack" + / "smoke_runs" + / "lattice_field" + / "seed_0" + / "lattice_eval.json", + pairwise_ranking_accuracy=0.01, + ) + + out = tmp_path / "report" + result = subprocess.run( + [ + sys.executable, + "scripts/report_hpc_clean_results.py", + "--inputs", + str(root), + "--out", + str(out), + ], + check=True, + cwd=Path(__file__).resolve().parents[1], + text=True, + capture_output=True, + ) + + manifest = json.loads((out / "clean_result_manifest.json").read_text(encoding="utf-8")) + assert manifest["num_rows"] == 4 + assert manifest["num_excluded"] == 2 + assert "report:" in result.stdout + + with (out / "clean_result_rows.csv").open("r", encoding="utf-8", newline="") as handle: + rows = list(csv.DictReader(handle)) + assert len(rows) == 4 + assert all("pilot" not in row["source_path"] for row in rows) + + report = (out / "clean_result_summary.md").read_text(encoding="utf-8") + assert "scaling_fixed14k_pick_common_eval" in report + assert "Excluded Paths" in report + + +def test_clean_hpc_report_names_fieldpreference_runs_separately(tmp_path: Path) -> None: + root = tmp_path / "experiments" + _write_eval( + root + / "maniskill_presuccess_six_task_runs" + / "legacy" + / "seed_0" + / "lattice_eval.json", + objective="legacy", + pairwise_ranking_accuracy=0.68, + selected_success_rate=0.31, + ) + _write_eval( + root + / "maniskill_presuccess_six_task_fieldpref_cpu_runs" + / "lattice_field" + / "seed_0" + / "lattice_eval.json", + objective="lattice_field", + pairwise_ranking_accuracy=0.74, + selected_success_rate=0.36, + ) + _write_eval( + root + / "maniskill_presuccess_six_task_visual_fieldpref_runs" + / "lattice_field" + / "seed_0" + / "lattice_eval.json", + objective="lattice_field", + observation_mode="rgb", + pairwise_ranking_accuracy=0.71, + selected_success_rate=0.33, + ) + _write_eval( + root + / "maniskill_presuccess_six_task_actionfix_cpu_runs" + / "lattice_field" + / "seed_0" + / "lattice_eval.json", + objective="lattice_field", + pairwise_ranking_accuracy=0.81, + selected_success_rate=0.34, + ) + _write_eval( + root + / "maniskill_presuccess_six_task_clip_actionfix_runs" + / "lattice_field" + / "seed_0" + / "lattice_eval.json", + objective="lattice_field", + observation_mode="rgb", + pairwise_ranking_accuracy=0.79, + ) + _write_eval( + root + / "maniskill_presuccess_six_task_rgb_actionfix_runs" + / "lattice_field" + / "seed_0" + / "lattice_eval.json", + objective="lattice_field", + observation_mode="rgb", + pairwise_ranking_accuracy=0.75, + ) + _write_eval( + root + / "maniskill_presuccess_transfer_leave_stack" + / "full_runs" + / "lattice_field" + / "seed_0" + / "lattice_eval.json", + objective="lattice_field", + pairwise_ranking_accuracy=0.62, + selected_success_rate=0.12, + ) + + out = tmp_path / "report" + subprocess.run( + [ + sys.executable, + "scripts/report_hpc_clean_results.py", + "--inputs", + str(root), + "--out", + str(out), + ], + check=True, + cwd=Path(__file__).resolve().parents[1], + text=True, + capture_output=True, + ) + + with (out / "clean_result_rows.csv").open("r", encoding="utf-8", newline="") as handle: + rows = list(csv.DictReader(handle)) + experiments = {row["experiment"] for row in rows} + assert "six_task_state" in experiments + assert "six_task_state_fieldpref" in experiments + assert "six_task_rgb_fieldpref" in experiments + assert "six_task_state_actionfix" in experiments + assert "six_task_rgb_clip_actionfix" in experiments + assert "six_task_rgb_actionfix" in experiments + assert "transfer_leave_stack_state" in experiments + + report = (out / "clean_result_summary.md").read_text(encoding="utf-8") + assert "six_task_state_fieldpref" in report + assert "six_task_rgb_fieldpref" in report + assert "six_task_state_actionfix" in report + assert "six_task_rgb_clip_actionfix" in report + assert "six_task_rgb_actionfix" in report + assert "transfer_leave_stack_state" in report + assert "Field-preference IAF pairwise ranking does not beat legacy" not in report + + +def test_clean_hpc_report_keeps_lattice_and_policy_rollout_separate(tmp_path: Path) -> None: + run = ( + tmp_path + / "experiments" + / "maniskill_presuccess_six_task_actionfix_cpu_runs" + / "lattice_field" + / "seed_2" + ) + _write_eval(run / "lattice_eval.json", seed=2) + (run / "policy_rollout.json").write_text( + json.dumps( + { + "objective": "lattice_field", + "observation_mode": "state", + "num_groups": 700, + "policy_rollout_success_rate": 0.30, + "policy_rollout_progress": 0.56, + "expert_success_rate": 0.37, + "oracle_success_rate": 0.43, + "policy_oracle_regret": 0.25, + "action_mse_to_best": 0.44, + } + ), + encoding="utf-8", + ) + + out = tmp_path / "report" + subprocess.run( + [ + sys.executable, + "scripts/report_hpc_clean_results.py", + "--inputs", + str(tmp_path / "experiments"), + "--out", + str(out), + ], + check=True, + cwd=Path(__file__).resolve().parents[1], + text=True, + capture_output=True, + ) + + with (out / "clean_result_rows.csv").open("r", encoding="utf-8", newline="") as handle: + rows = list(csv.DictReader(handle)) + assert {row["evaluation_kind"] for row in rows} == {"lattice", "policy_rollout"} + rollout = next(row for row in rows if row["evaluation_kind"] == "policy_rollout") + assert rollout["seed"] == "2" + assert rollout["policy_rollout_success_rate"] == "0.3" + + +def test_clean_hpc_report_warns_on_weak_heldout_transfer(tmp_path: Path) -> None: + root = tmp_path / "experiments" + _write_eval( + root + / "maniskill_presuccess_transfer_leave_stack" + / "state_actionfix_runs" + / "lattice_field" + / "seed_0" + / "lattice_eval.json", + selected_success_rate=0.02, + ) + + out = tmp_path / "report" + subprocess.run( + [ + sys.executable, + "scripts/report_hpc_clean_results.py", + "--inputs", + str(root), + "--out", + str(out), + ], + check=True, + cwd=Path(__file__).resolve().parents[1], + text=True, + capture_output=True, + ) + + report = (out / "clean_result_summary.md").read_text(encoding="utf-8") + assert "do not claim broad OOD task transfer" in report diff --git a/tests/test_intervention_sampler.py b/tests/test_intervention_sampler.py new file mode 100644 index 0000000000000000000000000000000000000000..8addc3c75cc98ab9bc55096ccb8b97d24f33de3c --- /dev/null +++ b/tests/test_intervention_sampler.py @@ -0,0 +1,93 @@ +from __future__ import annotations + +from dovla_cil.interventions.samplers import InterventionSampler, RandomInterventionSampler +from dovla_cil.sim.registry import get_simulator_backend +from dovla_cil.tasks.library import ToyTaskLibrary + + +def test_sampler_returns_exactly_k_or_less_when_impossible() -> None: + task = ToyTaskLibrary().get_by_id("toy_pick_object_among_distractors") + sim = get_simulator_backend("toy") + sim.seed(3) + sim.reset_task(task) + expert = RandomInterventionSampler(k=1, seed=3).sample(sim.render_observation(), task)[0].action + + actions = InterventionSampler(seed=9).sample( + task=task, + observation=sim.render_observation(), + symbolic_state=sim.get_symbolic_state(), + expert_actions=[expert], + k=6, + ) + + assert len(actions) == 6 + assert len({action.action_id for action in actions}) == len(actions) + + +def test_sampler_candidate_metadata_and_determinism() -> None: + task = ToyTaskLibrary().get_by_id("toy_put_red_mug_in_blue_bowl") + sim = get_simulator_backend("toy") + sim.seed(5) + sim.reset_task(task) + expert = RandomInterventionSampler(k=1, seed=5).sample(sim.render_observation(), task)[0].action + + first = InterventionSampler(seed=11).sample( + task, sim.render_observation(), sim.get_symbolic_state(), [expert], 8 + ) + second = InterventionSampler(seed=11).sample( + task, sim.render_observation(), sim.get_symbolic_state(), [expert], 8 + ) + + assert [action.to_dict() for action in first] == [action.to_dict() for action in second] + for action in first: + assert "candidate_type" in action.metadata + assert "intended_target" in action.metadata + assert "intended_relation" in action.metadata + assert "difficulty" in action.metadata + + +def test_wrong_target_actions_use_distractors_if_present() -> None: + task = ToyTaskLibrary().get_by_id("toy_pick_object_among_distractors") + sim = get_simulator_backend("toy") + sim.seed(7) + sim.reset_task(task) + + actions = InterventionSampler(seed=7).sample( + task, + sim.render_observation(), + sim.get_symbolic_state(), + [], + 8, + ) + wrong_targets = [ + action.metadata["intended_target"] + for action in actions + if action.metadata.get("candidate_type") == "wrong_target" + ] + + assert wrong_targets + assert set(wrong_targets).issubset(set(task.distractor_object_ids)) + + +def test_near_miss_actions_differ_from_expert() -> None: + task = ToyTaskLibrary().get_by_id("toy_put_red_mug_in_blue_bowl") + sim = get_simulator_backend("toy") + sim.seed(13) + sim.reset_task(task) + expert = RandomInterventionSampler(k=1, seed=13).sample( + sim.render_observation(), task + )[0].action + + actions = InterventionSampler(seed=13).sample( + task, + sim.render_observation(), + sim.get_symbolic_state(), + [expert], + 8, + ) + near_misses = [ + action for action in actions if action.metadata.get("candidate_type") == "near_miss" + ] + + assert near_misses + assert all(action.values != expert.values for action in near_misses) diff --git a/tests/test_lattice_eval.py b/tests/test_lattice_eval.py new file mode 100644 index 0000000000000000000000000000000000000000..c837bd5ceb84201bc4112d52fa10687e5f31e5f2 --- /dev/null +++ b/tests/test_lattice_eval.py @@ -0,0 +1,34 @@ +from __future__ import annotations + +from dovla_cil.eval.lattice_eval import ( + _empty_task_stats, + _finalize_task_stats, + _ndcg, + _validation_group_ids, +) + + +def test_lattice_ndcg_rewards_correct_order() -> None: + utilities = [2.0, 1.0, 0.0] + + assert _ndcg([3.0, 2.0, 1.0], utilities) == 1.0 + assert _ndcg([1.0, 2.0, 3.0], utilities) < 1.0 + + +def test_lattice_validation_split_is_group_deterministic() -> None: + groups = [f"g{index}" for index in range(10)] + + first = _validation_group_ids(groups, val_fraction=0.2, seed=7) + second = _validation_group_ids(groups, val_fraction=0.2, seed=7) + + assert first == second + assert len(first) == 2 + assert len(set(first)) == len(first) + + +def test_edge_metrics_are_undefined_without_comparable_pairs() -> None: + metrics = _finalize_task_stats(_empty_task_stats()) + + assert metrics["num_pairs"] == 0 + assert metrics["pairwise_ranking_accuracy"] is None + assert metrics["potential_edge_mae"] is None diff --git a/tests/test_lerobot_export.py b/tests/test_lerobot_export.py new file mode 100644 index 0000000000000000000000000000000000000000..0f2e6b2945379057affaeec9a0473d2eaab1386f --- /dev/null +++ b/tests/test_lerobot_export.py @@ -0,0 +1,107 @@ +from __future__ import annotations + +import subprocess +import sys +from pathlib import Path + +from dovla_cil.data.datasets import CILDataset +from dovla_cil.data.lerobot_export import LeRobotExportConfig, export_lerobot_style_dataset +from dovla_cil.generation.pipeline import generate_cil_dataset +from dovla_cil.tasks.library import built_in_toy_tasks +from dovla_cil.utils.io import iter_jsonl, read_json + + +def _make_toy_cil(tmp_path: Path) -> Path: + dataset_dir = tmp_path / "cil" + generate_cil_dataset( + backend="toy", + tasks=built_in_toy_tasks()[:2], + out_dir=dataset_dir, + num_states_per_task=2, + k=4, + seed=19, + shard_size=8, + inline_observations=True, + ) + return dataset_dir + + +def test_lerobot_style_export_selects_best_record_per_group(tmp_path: Path) -> None: + dataset_dir = _make_toy_cil(tmp_path) + out_dir = tmp_path / "lerobot" + + metadata = export_lerobot_style_dataset( + LeRobotExportConfig( + dataset_dir=dataset_dir, + out_dir=out_dir, + max_groups=3, + copy_images=False, + ) + ) + + rows = list(iter_jsonl(out_dir / "train.jsonl")) + dataset = CILDataset(dataset_dir) + + assert metadata["schema_version"] == "dovla-cil-lerobot-export/v0" + assert metadata["num_episodes"] == 3 + assert len(rows) == 3 + assert (out_dir / "tasks.jsonl").exists() + assert read_json(out_dir / "metadata.json") == metadata + for row in rows: + group = dataset.get_group(row["cil"]["group_id"]) + assert row["reward"] == max(record.reward.score for record in group) + assert row["cil"]["record_id"] in {record.record_id for record in group} + assert row["task"] + assert row["observation"]["image"] is None + assert "action_chunk" in row + + +def test_lerobot_style_export_cli_runs_without_network(tmp_path: Path) -> None: + dataset_dir = _make_toy_cil(tmp_path) + out_dir = tmp_path / "cli-export" + + result = subprocess.run( + [ + sys.executable, + "scripts/export_lerobot_dataset.py", + "--dataset", + str(dataset_dir), + "--out", + str(out_dir), + "--max-groups", + "2", + "--no-images", + ], + check=True, + capture_output=True, + text=True, + ) + + assert "dovla-cil-lerobot-export/v0" in result.stdout + assert len(list(iter_jsonl(out_dir / "train.jsonl"))) == 2 + + +def test_task_balanced_export_covers_tasks_deterministically(tmp_path: Path) -> None: + dataset_dir = _make_toy_cil(tmp_path) + first_out = tmp_path / "balanced-first" + second_out = tmp_path / "balanced-second" + config_kwargs = { + "dataset_dir": dataset_dir, + "max_groups": 4, + "group_sampling": "task_balanced", + "seed": 7, + "copy_images": False, + } + + export_lerobot_style_dataset(LeRobotExportConfig(out_dir=first_out, **config_kwargs)) + export_lerobot_style_dataset(LeRobotExportConfig(out_dir=second_out, **config_kwargs)) + first = list(iter_jsonl(first_out / "train.jsonl")) + second = list(iter_jsonl(second_out / "train.jsonl")) + + assert {row["cil"]["task_id"] for row in first} == { + "toy_pick_red_mug", + "toy_put_red_mug_in_blue_bowl", + } + assert [row["cil"]["group_id"] for row in first] == [ + row["cil"]["group_id"] for row in second + ] diff --git a/tests/test_losses.py b/tests/test_losses.py new file mode 100644 index 0000000000000000000000000000000000000000..3feb808a43e36e05067c1f7df66b68d74d67202c --- /dev/null +++ b/tests/test_losses.py @@ -0,0 +1,202 @@ +from __future__ import annotations + +import pytest + +from dovla_cil.training.losses import ( + CompositeLoss, + behavior_cloning_loss, + causal_contrastive_loss, + effect_prediction_loss, + language_minimal_pair_loss, + lattice_cycle_residual, + lattice_field_loss, + pairwise_ranking_loss, + progress_loss, + regret_loss, + regret_targets, + same_state_pairwise_ranking_loss, + success_loss, +) + + +def test_ranking_loss_prefers_correct_order() -> None: + rewards = [1.0, 0.0, -1.0] + good_scores = [2.0, 1.0, 0.0] + bad_scores = [0.0, 1.0, 2.0] + good_loss = same_state_pairwise_ranking_loss(good_scores, rewards) + bad_loss = same_state_pairwise_ranking_loss(bad_scores, rewards) + assert good_loss < bad_loss + + +def test_ranking_loss_ignores_ties() -> None: + assert same_state_pairwise_ranking_loss([0.0, 1.0], [1.0, 1.0]) == 0.0 + + +def test_ranking_loss_checks_shape() -> None: + with pytest.raises(ValueError): + same_state_pairwise_ranking_loss([0.0], [0.0, 1.0]) + + +def test_regret_targets() -> None: + assert regret_targets([1.0, 0.25, -1.0]) == [0.0, 0.75, 2.0] + + +def test_pairwise_ranking_loss_lower_when_order_is_correct() -> None: + good = pairwise_ranking_loss([2.0, 1.0], [0.0, 0.0], [1.0, 1.0], [0.0, -1.0]) + bad = pairwise_ranking_loss([0.0, 0.0], [2.0, 1.0], [1.0, 1.0], [0.0, -1.0]) + assert good < bad + + +def test_regret_loss_zero_when_exact() -> None: + assert regret_loss([0.0, 0.5, 1.0], [0.0, 0.5, 1.0]) == 0.0 + + +def test_behavior_cloning_loss_works() -> None: + assert behavior_cloning_loss([1.0, 2.0], [1.0, 4.0]) == pytest.approx(2.0) + + +def test_effect_prediction_loss_combines_continuous_and_binary_terms() -> None: + loss = effect_prediction_loss( + {"continuous": [0.0, 1.0], "binary_logits": [0.0]}, + {"continuous": [0.0, 3.0], "binary": [1.0]}, + ) + assert loss > 0.0 + + +def test_success_loss_prefers_correct_logits() -> None: + assert success_loss([3.0], [1.0]) < success_loss([-3.0], [1.0]) + + +def test_progress_loss_works() -> None: + assert progress_loss([0.5], [0.5]) == pytest.approx(0.0) + + +def test_contrastive_loss_finite() -> None: + loss = causal_contrastive_loss( + [[1.0, 0.0]], + [[1.0, 0.0]], + [[0.0, 1.0]], + temperature=0.1, + ) + assert float(loss) >= 0.0 + + +def test_language_minimal_pair_loss_pushes_and_pulls() -> None: + close_different = language_minimal_pair_loss([[0.0, 0.0]], [[0.1, 0.0]], [True], margin=1.0) + far_different = language_minimal_pair_loss([[0.0, 0.0]], [[2.0, 0.0]], [True], margin=1.0) + same_identical = language_minimal_pair_loss([[0.0, 0.0]], [[0.0, 0.0]], [False], margin=1.0) + same_apart = language_minimal_pair_loss([[0.0, 0.0]], [[1.0, 0.0]], [False], margin=1.0) + assert far_different < close_different + assert same_identical < same_apart + + +def test_composite_returns_components_and_total() -> None: + output = CompositeLoss()( + predictions={ + "pred_action": [1.0, 2.0], + "pred_regret": [0.0, 1.0], + "pred_scores_i": [2.0], + "pred_scores_j": [0.0], + }, + targets={ + "target_action": [1.0, 3.0], + "target_regret": [0.0, 1.0], + "rewards_i": [1.0], + "rewards_j": [0.0], + }, + ) + assert set(output) >= {"total", "bc", "rank", "regret"} + assert float(output["total"]) >= 0.0 + + +def test_lattice_field_loss_is_invariant_to_state_reward_offsets() -> None: + torch = pytest.importorskip("torch") + potential = torch.tensor([0.2, -0.1, 0.7, 0.0]) + utility = torch.tensor([0.8, 0.3, 0.6, 0.1]) + effect = torch.tensor([[0.0], [0.2], [0.5], [0.1]]) + target_effect = torch.tensor([[0.1], [0.4], [0.6], [0.0]]) + group_ids = ["state-a", "state-a", "state-b", "state-b"] + + base = lattice_field_loss(potential, utility, effect, target_effect, group_ids) + shifted = lattice_field_loss( + potential, + utility + torch.tensor([17.0, 17.0, -9.0, -9.0]), + effect, + target_effect, + group_ids, + ) + + assert torch.allclose(base["potential"], shifted["potential"]) + assert base["edge_count"] == shifted["edge_count"] == 2 + + +def test_lattice_field_is_zero_under_groupwise_gauge_shifts() -> None: + torch = pytest.importorskip("torch") + utility = torch.tensor([0.8, 0.3, 0.6, 0.1]) + target_effect = torch.tensor([[0.1, 0.2], [0.4, 0.0], [0.6, -0.2], [0.0, 0.3]]) + group_ids = ["state-a", "state-a", "state-b", "state-b"] + potential = utility + torch.tensor([5.0, 5.0, -2.0, -2.0]) + predicted_effect = target_effect + torch.tensor( + [[1.0, -3.0], [1.0, -3.0], [-4.0, 2.0], [-4.0, 2.0]] + ) + + loss = lattice_field_loss( + potential, + utility, + predicted_effect, + target_effect, + group_ids, + ) + + assert float(loss["potential"]) == pytest.approx(0.0, abs=1e-7) + assert float(loss["effect"]) == pytest.approx(0.0, abs=1e-7) + + +def test_lattice_field_orientation_penalizes_reversed_edge_order() -> None: + torch = pytest.importorskip("torch") + utility = torch.tensor([1.0, 0.0]) + effect = torch.zeros((2, 2)) + correct = lattice_field_loss( + torch.tensor([1.0, 0.0]), + utility, + effect, + effect, + ["state", "state"], + ) + reversed_order = lattice_field_loss( + torch.tensor([0.0, 1.0]), + utility, + effect, + effect, + ["state", "state"], + ) + + assert float(correct["potential"]) == pytest.approx(0.0, abs=1e-7) + assert float(reversed_order["orientation"]) > 0.0 + assert float(reversed_order["potential"]) > float(correct["potential"]) + assert float(reversed_order["preference"]) > float(correct["preference"]) + + +def test_lattice_field_preference_is_group_offset_invariant() -> None: + torch = pytest.importorskip("torch") + potential = torch.tensor([0.4, -0.2, 2.0, 1.7]) + utility = torch.tensor([0.9, 0.1, 0.8, 0.2]) + effect = torch.zeros((4, 2)) + group_ids = ["a", "a", "b", "b"] + + base = lattice_field_loss(potential, utility, effect, effect, group_ids) + shifted = lattice_field_loss( + potential, + utility + torch.tensor([10.0, 10.0, -4.0, -4.0]), + effect, + effect, + group_ids, + ) + + assert torch.allclose(base["preference"], shifted["preference"]) + + +def test_scalar_potential_has_zero_cycle_residual() -> None: + torch = pytest.importorskip("torch") + residual = lattice_cycle_residual(torch.tensor([0.3, -1.2, 2.4]), [[0, 1, 2]]) + assert float(residual) == pytest.approx(0.0, abs=1e-7) diff --git a/tests/test_manifest_runner.py b/tests/test_manifest_runner.py new file mode 100644 index 0000000000000000000000000000000000000000..37114e4f4e5666e3bb4a3a1df4db6ec7f01a62f0 --- /dev/null +++ b/tests/test_manifest_runner.py @@ -0,0 +1,229 @@ +from __future__ import annotations + +import json +import os +import subprocess +import sys +from pathlib import Path + +import pytest + +from scripts.run_manifest import PlannedJob, execute_local_jobs, load_manifest, plan_jobs +from scripts.train_dovla import _parse_loss_weights + + +def test_manifest_dry_run_writes_resolved_manifest(tmp_path: Path) -> None: + out_dir = tmp_path / "run" + result = subprocess.run( + [ + sys.executable, + "scripts/run_manifest.py", + "manifests/scaling_k_sweep.yaml", + "--dry-run", + "--out", + str(out_dir), + ], + check=True, + capture_output=True, + text=True, + ) + + assert "manifest: scaling_k_sweep" in result.stdout + assert "planned jobs:" in result.stdout + assert (out_dir / "resolved_manifest.yaml").exists() + jobs_path = out_dir / "planned_jobs.json" + assert jobs_path.exists() + jobs = json.loads(jobs_path.read_text(encoding="utf-8")) + assert any(job["stage"] == "dataset_generation" for job in jobs) + assert any(job["stage"] == "training" for job in jobs) + assert any(job["stage"] == "evaluation" for job in jobs) + assert any(job["stage"] == "scaling_sweeps" for job in jobs) + generation_job = next(job for job in jobs if job["stage"] == "dataset_generation") + assert "--num-tasks" in generation_job["command"] + + +def test_manifest_redacts_secrets_from_outputs(tmp_path: Path) -> None: + secret = "manifest_secret_value" + manifest = tmp_path / "manifest.yaml" + manifest.write_text( + """ +name: secret_test +run_dir: ${TEST_RUN_DIR:-runs/secret_test} +dataset_generation: + backend: toy + simulator_params: {} + task_source: builtins + num_tasks: 1 + num_states_per_task: 1 + k: 2 + shard_size: 8 + output_path: outputs/secret_test/cil + seed: 0 +vlm_annotation: + enabled: false + cache_path: outputs/secret_test/cache.json + model_env_var: OPENCLAUDE_MODEL + api_key: ${OPENCLAUDE_API_KEY} +training: + model_size: tiny + hidden_dim: 32 + batch_groups: 1 + records_per_group: 2 + learning_rate: 0.001 + loss_weights: {bc: 1.0} + epochs: 1 + steps: null + checkpoint_path: outputs/secret_test/train/best.pt +evaluation: + causalstress: + enabled: true + backend: toy + num_tasks: 1 + k: 2 + output_path: outputs/secret_test/eval/causalstress.json + libero: {enabled: false, placeholder: true} + maniskill: {enabled: false, placeholder: true} + simpler: {enabled: false, placeholder: true} +baselines: + enabled: false + output_root: outputs/secret_test/baselines + names: [] +scaling_sweeps: + enabled: false + output_path: outputs/secret_test/scaling + total_records: 4 + k_values: [1, 2] + epochs: 1 +""", + encoding="utf-8", + ) + out_dir = tmp_path / "run" + env = {**os.environ, "OPENCLAUDE_API_KEY": secret, "TEST_RUN_DIR": str(out_dir)} + + result = subprocess.run( + [ + sys.executable, + "scripts/run_manifest.py", + str(manifest), + "--dry-run", + "--emit-slurm", + "--out", + str(out_dir), + ], + check=True, + capture_output=True, + text=True, + env=env, + ) + + all_text = result.stdout + result.stderr + all_text += (out_dir / "resolved_manifest.yaml").read_text(encoding="utf-8") + all_text += (out_dir / "planned_jobs.json").read_text(encoding="utf-8") + for path in (out_dir / "slurm").glob("*.sbatch"): + all_text += path.read_text(encoding="utf-8") + assert secret not in all_text + assert "" in (out_dir / "resolved_manifest.yaml").read_text(encoding="utf-8") + + +def test_manifest_files_exist() -> None: + for path in ( + "manifests/cil_160m.yaml", + "manifests/cil_1b_template.yaml", + "manifests/scaling_k_sweep.yaml", + "manifests/baselines_full.yaml", + ): + assert Path(path).exists() + + +def test_large_manifests_plan_measured_maniskill_generation() -> None: + expected_budgets = { + "manifests/cil_160m.yaml": 160_000_000, + "manifests/cil_1b_template.yaml": 1_000_000_000, + } + for path, expected_budget in expected_budgets.items(): + manifest = load_manifest(path) + generation = manifest["dataset_generation"] + assert ( + generation["num_tasks"] + * generation["num_states_per_task"] + * generation["k"] + == expected_budget + ) + job = next(job for job in plan_jobs(manifest) if job.stage == "dataset_generation") + assert job.command[1] == "scripts/generate_maniskill_lattice.py" + assert "--demo" in job.command + assert "--state-batch-size" in job.command + assert "--parallel-branches" in job.command + groups_index = job.command.index("--num-groups") + 1 + assert int(job.command[groups_index]) * generation["k"] == expected_budget + evaluation_job = next( + item for item in plan_jobs(manifest) if item.name == "eval_causalstress" + ) + assert not evaluation_job.local_executable + + +def test_training_plan_carries_manifest_loss_weights() -> None: + manifest = load_manifest("manifests/scaling_k_sweep.yaml") + job = next(job for job in plan_jobs(manifest) if job.stage == "training") + encoded = [ + job.command[index + 1] + for index, value in enumerate(job.command[:-1]) + if value == "--loss-weight" + ] + assert "bc=1.0" in encoded + assert "rank=1.0" in encoded + parsed = _parse_loss_weights(encoded) + assert parsed.bc == 1.0 + assert parsed.rank == 1.0 + + +def test_manifest_validation_rejects_invalid_budget(tmp_path: Path) -> None: + payload = Path("manifests/scaling_k_sweep.yaml").read_text(encoding="utf-8") + invalid = tmp_path / "invalid.yaml" + invalid.write_text(payload.replace("num_tasks: 10", "num_tasks: 0"), encoding="utf-8") + + with pytest.raises(ValueError, match="num_tasks"): + load_manifest(invalid) + + +def test_emitted_slurm_directives_are_concrete(tmp_path: Path) -> None: + out_dir = tmp_path / "run" + subprocess.run( + [ + sys.executable, + "scripts/run_manifest.py", + "manifests/scaling_k_sweep.yaml", + "--dry-run", + "--emit-slurm", + "--out", + str(out_dir), + ], + check=True, + capture_output=True, + text=True, + ) + + scripts = list((out_dir / "slurm").glob("*.sbatch")) + assert scripts + for script in scripts: + text = script.read_text(encoding="utf-8") + directives = "\n".join(line for line in text.splitlines() if line.startswith("#SBATCH")) + assert "${" not in directives + + +def test_local_execution_uses_current_interpreter(tmp_path: Path) -> None: + result_path = tmp_path / "interpreter.txt" + job = PlannedJob( + name="interpreter_probe", + stage="test", + command=[ + "python", + "-c", + f"import pathlib,sys; pathlib.Path({str(result_path)!r}).write_text(sys.executable)", + ], + local_executable=True, + ) + + execute_local_jobs([job]) + + assert Path(result_path.read_text(encoding="utf-8")).resolve() == Path(sys.executable).resolve() diff --git a/tests/test_maniskill_backend.py b/tests/test_maniskill_backend.py new file mode 100644 index 0000000000000000000000000000000000000000..8e34d2c24b57315e483dfd652b202ccc882ef1b4 --- /dev/null +++ b/tests/test_maniskill_backend.py @@ -0,0 +1,62 @@ +from __future__ import annotations + +import pytest + +from dovla_cil.sim import genesis_backend, maniskill_backend +from dovla_cil.sim.registry import create_backend, get_backend_class, list_backends +from dovla_cil.sim.toy_backend import ToyBackend + + +def test_registry_can_list_optional_backends() -> None: + backends = list_backends() + + assert "toy" in backends + assert "maniskill" in backends + assert "genesis" in backends + + +def test_maniskill_backend_class_imports_without_dependency() -> None: + backend_class = get_backend_class("maniskill") + + assert backend_class.__name__ == "ManiSkillBackend" + + +def test_genesis_backend_class_imports_without_dependency() -> None: + backend_class = get_backend_class("genesis") + + assert backend_class.__name__ == "GenesisBackend" + + +def test_requesting_maniskill_without_dependency_gives_clean_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(maniskill_backend, "_find_maniskill_module", lambda: None) + + with pytest.raises(ImportError) as exc_info: + create_backend("maniskill") + + message = str(exc_info.value) + assert "Install optional dependency with" in message + assert "OPENCLAUDE" not in message + + +def test_requesting_genesis_without_dependency_gives_clean_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(genesis_backend, "_find_genesis_module", lambda: None) + + with pytest.raises(ImportError) as exc_info: + create_backend("genesis") + + message = str(exc_info.value) + assert "Install optional dependency with" in message + assert "OPENCLAUDE" not in message + + +def test_unknown_backend_raises_value_error() -> None: + with pytest.raises(ValueError): + create_backend("unknown") + + +def test_toy_backend_unaffected() -> None: + assert isinstance(create_backend("toy"), ToyBackend) diff --git a/tests/test_maniskill_lattice.py b/tests/test_maniskill_lattice.py new file mode 100644 index 0000000000000000000000000000000000000000..e38530155de48f98b91fdf1a7689044d326d1b66 --- /dev/null +++ b/tests/test_maniskill_lattice.py @@ -0,0 +1,319 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import numpy as np +import torch + +from dovla_cil.generation.maniskill_lattice import ( + ManiSkillLatticeConfig, + _extract_rgb, + _lattice_quality_summary, + _trim_padded_branch_groups, + flatten_state, + get_maniskill_task_profile, + plan_branch_points, + sample_action_lattice, +) +from dovla_cil.generation.maniskill_parallel import ( + execute_action_lattice_batch, + execute_grouped_action_lattice_batch, + repeat_state_batch, + slice_state_batch, +) +from dovla_cil.models.dovla import vectorize_toy_observation + + +def test_maniskill_action_lattice_is_deterministic_and_keeps_expert() -> None: + expert = np.full((4, 7), 0.25, dtype=np.float32) + first = sample_action_lattice(expert, k=8, rng=np.random.default_rng(7)) + second = sample_action_lattice(expert, k=8, rng=np.random.default_rng(7)) + + assert len(first) == 8 + assert first[0]["candidate_type"] == "expert" + assert np.array_equal(first[0]["values"], expert) + assert [item["candidate_type"] for item in first] == [ + item["candidate_type"] for item in second + ] + assert all( + np.array_equal(left["values"], right["values"]) + for left, right in zip(first, second, strict=True) + ) + assert any(not np.array_equal(item["values"], expert) for item in first[1:]) + + +def test_random_negative_mode_changes_actions_not_only_labels() -> None: + expert = np.full((4, 7), 0.25, dtype=np.float32) + + candidates = sample_action_lattice( + expert, + k=8, + rng=np.random.default_rng(17), + mode="random", + ) + + assert candidates[0]["candidate_type"] == "expert" + assert all(item["candidate_type"] == "random_negative" for item in candidates[1:]) + assert all(not np.array_equal(item["values"], expert) for item in candidates[1:]) + assert all(item["perturbation"]["kind"] == "full_range_uniform" for item in candidates[1:]) + + +def test_maniskill_task_profiles_cover_downloaded_multitask_suite() -> None: + expected_targets = { + "PickCube-v1": ("cube",), + "PushCube-v1": ("cube",), + "PullCube-v1": ("cube",), + "StackCube-v1": ("cubeA",), + "PegInsertionSide-v1": ("peg",), + "LiftPegUpright-v1": ("peg",), + } + + for env_id, target_actors in expected_targets.items(): + profile = get_maniskill_task_profile(env_id) + assert profile.target_actors == target_actors + assert set(profile.target_actors).issubset(profile.effect_actors) + assert set(profile.reference_actors).issubset(profile.effect_actors) + assert profile.instruction + assert profile.skill_type + + +def test_unknown_maniskill_task_profile_fails_before_simulation() -> None: + try: + get_maniskill_task_profile("UnknownTask-v0") + except ValueError as exc: + assert "Unsupported ManiSkill lattice task" in str(exc) + else: # pragma: no cover - protects explicit task semantics + raise AssertionError("unknown tasks must not silently use PickCube semantics") + + +def test_rgb_extraction_uses_batched_camera_tensor() -> None: + expected = torch.arange(2 * 4 * 5 * 3, dtype=torch.uint8).reshape(2, 4, 5, 3) + observation = { + "sensor_data": { + "base_camera": {"rgb": expected}, + } + } + + actual = _extract_rgb(observation) + + assert actual.shape == (2, 4, 5, 3) + assert actual.dtype == np.uint8 + assert np.array_equal(actual, expected.numpy()) + + +def test_rgb_generation_requires_parallel_branches(tmp_path) -> None: + try: + ManiSkillLatticeConfig( + demo_path=tmp_path / "demo.h5", + output_dir=tmp_path / "out", + obs_mode="state+rgb", + parallel_branches=False, + ) + except ValueError as exc: + assert "RGB capture requires parallel_branches" in str(exc) + else: # pragma: no cover - prevents silently losing K final images + raise AssertionError("serial RGB generation must be rejected") + + +def test_branch_plan_is_duplicate_free_and_stable_across_worker_slices() -> None: + lengths = {"traj_0": 5, "traj_1": 6, "traj_2": 2} + full = plan_branch_points(lengths, horizon=3, seed=11) + repeated = plan_branch_points(lengths, horizon=3, seed=11) + + assert full == repeated + assert len(full) == 7 + assert len(set(full)) == len(full) + assert full[:3] + full[3:] == full + assert all(name != "traj_2" for name, _step in full) + + +def test_branch_plan_excludes_states_that_are_already_successful() -> None: + lengths = {"traj_0": 6, "traj_1": 5} + success = { + "traj_0": np.asarray([False, False, True, True, True, True]), + "traj_1": np.asarray([False, False, False, False, True]), + } + + plan = plan_branch_points( + lengths, + horizon=2, + seed=3, + success_flags=success, + ) + + assert set(plan) == { + ("traj_0", 0), + ("traj_0", 1), + ("traj_0", 2), + ("traj_1", 0), + ("traj_1", 1), + ("traj_1", 2), + ("traj_1", 3), + } + assert ("traj_0", 3) not in plan + assert ("traj_0", 4) not in plan + + +def test_lattice_quality_summary_detects_degenerate_groups() -> None: + def record(group: str, candidate: str, progress: float, success: bool): + return SimpleNamespace( + group_id=group, + candidate_type=candidate, + reward=SimpleNamespace(progress=progress, success=success), + ) + + summary = _lattice_quality_summary( + [ + record("g0", "expert", 1.0, True), + record("g0", "no_op", 0.0, False), + record("g1", "expert", 0.5, False), + record("g1", "no_op", 0.5, False), + ] + ) + + assert summary["expert_success_rate"] == 0.5 + assert summary["no_op_success_rate"] == 0.0 + assert summary["mean_reward_spread"] == 0.5 + assert summary["nondegenerate_group_fraction"] == 0.5 + + +def test_partial_state_batch_drops_only_padding_results() -> None: + branch_groups = [["g0"], ["g1"], ["g2"], ["padding"]] + + actual = _trim_padded_branch_groups(branch_groups, real_count=3) + + assert actual == [["g0"], ["g1"], ["g2"]] + + +def test_maniskill_state_features_feed_core_observation_vectorizer() -> None: + state = { + "actors": {"cube": np.asarray([[1.0, 2.0, 3.0]], dtype=np.float32)}, + "articulations": {"panda": np.asarray([[4.0, 5.0]], dtype=np.float32)}, + } + features = flatten_state(state) + vector = vectorize_toy_observation({"features": features}, obs_dim=8) + + assert features == [1.0, 2.0, 3.0, 4.0, 5.0] + assert vector[:5] == features + assert len(vector) == 8 + + +def test_maniskill_state_can_be_broadcast_to_identical_parallel_branches() -> None: + state = { + "actors": {"cube": np.asarray([[1.0, 2.0, 3.0]], dtype=np.float32)}, + "articulations": {"panda": np.asarray([[4.0, 5.0]], dtype=np.float32)}, + } + + repeated = repeat_state_batch(state, 4) + + assert repeated["actors"]["cube"].shape == (4, 3) + assert np.array_equal(repeated["actors"]["cube"][0], repeated["actors"]["cube"][3]) + assert np.array_equal(slice_state_batch(repeated, 2)["articulations"]["panda"], [[4.0, 5.0]]) + + +def test_maniskill_parallel_state_rejects_non_singleton_source() -> None: + state = {"actors": {"cube": np.zeros((2, 3), dtype=np.float32)}} + + try: + repeat_state_batch(state, 4) + except ValueError as exc: + assert "exactly one environment" in str(exc) + else: # pragma: no cover - protects the exact-state invariant + raise AssertionError("non-singleton source state should be rejected") + + +def test_parallel_lattice_executes_distinct_branches_from_repeated_state() -> None: + class Controller: + reset_count = 0 + + def reset(self) -> None: + self.reset_count += 1 + + class Agent: + controller = Controller() + + class FakeEnv: + agent = Agent() + + def set_state_dict(self, state): + self.state = state + + def step(self, action): + self.state["actors"]["cube"][:, 0] += action[:, 0] + reward = self.state["actors"]["cube"][:, 0] + return None, reward, None, None, {"success": reward > 0.5} + + def get_state_dict(self): + return self.state + + actions = np.asarray( + [ + [[0.1, 0.0], [0.2, 0.0]], + [[0.4, 0.0], [0.3, 0.0]], + ], + dtype=np.float32, + ) + state = {"actors": {"cube": np.zeros((1, 3), dtype=np.float32)}} + env = FakeEnv() + + after, rewards, successes, restore_error = execute_action_lattice_batch( + env, + state, + actions, + torch=torch, + device=torch.device("cpu"), + ) + + assert env.agent.controller.reset_count == 1 + assert restore_error == 0.0 + assert np.allclose(after["actors"]["cube"][:, 0], [0.3, 0.7]) + assert np.allclose(rewards, [0.3, 0.7]) + assert successes.tolist() == [False, True] + + +def test_parallel_lattice_keeps_state_groups_isolated() -> None: + class Controller: + def reset(self) -> None: + pass + + class Agent: + controller = Controller() + + class FakeEnv: + agent = Agent() + + def set_state_dict(self, state): + self.state = state + + def step(self, action): + self.state["actors"]["cube"][:, 0] += action[:, 0] + reward = self.state["actors"]["cube"][:, 0] + return None, reward, None, None, {"success": reward > 5.0} + + def get_state_dict(self): + return self.state + + states = [ + {"actors": {"cube": np.asarray([[0.0, 0.0]], dtype=np.float32)}}, + {"actors": {"cube": np.asarray([[10.0, 0.0]], dtype=np.float32)}}, + ] + actions = np.asarray( + [ + [[[1.0]], [[2.0]]], + [[[-1.0]], [[-2.0]]], + ], + dtype=np.float32, + ) + + after, rewards, successes, restore_error = execute_grouped_action_lattice_batch( + FakeEnv(), + states, + actions, + torch=torch, + device=torch.device("cpu"), + ) + + assert restore_error == 0.0 + assert np.allclose(after["actors"]["cube"][:, 0], [1.0, 2.0, 9.0, 8.0]) + assert np.allclose(rewards, [[1.0, 2.0], [9.0, 8.0]]) + assert successes.tolist() == [[False, False], [True, True]] diff --git a/tests/test_maniskill_policy_rollout.py b/tests/test_maniskill_policy_rollout.py new file mode 100644 index 0000000000000000000000000000000000000000..481473939d07f9bccef8315e63af08b2971b92b2 --- /dev/null +++ b/tests/test_maniskill_policy_rollout.py @@ -0,0 +1,113 @@ +from __future__ import annotations + +import pickle +from pathlib import Path + +import numpy as np + +from dovla_cil.data.schema import ( + CIL_VERSION, + ActionChunk, + CILRecord, + FailureInfo, + RewardInfo, + StructuredEffect, +) +from dovla_cil.eval.maniskill_policy_rollout import ( + _adapt_action_dim, + _load_state_archive, + _numeric_action_values, + _summarize_rows, +) + + +def test_policy_rollout_action_dim_adapter_slices_and_pads() -> None: + actions = np.ones((2, 3, 8), dtype=np.float32) + + sliced = _adapt_action_dim(actions, 7) + padded = _adapt_action_dim(actions[:, :, :6], 7) + + assert sliced.shape == (2, 3, 7) + assert np.allclose(sliced, 1.0) + assert padded.shape == (2, 3, 7) + assert np.allclose(padded[:, :, :6], 1.0) + assert np.allclose(padded[:, :, 6], 0.0) + + +def test_policy_rollout_summary_uses_measured_rollout_rows() -> None: + rows = [ + { + "success": True, + "progress": 0.8, + "oracle_success": True, + "expert_success": False, + "oracle_regret": 0.2, + "expert_regret": 0.0, + "action_mse_to_best": 0.1, + "restore_error": 1e-7, + }, + { + "success": False, + "progress": 0.2, + "oracle_success": True, + "expert_success": True, + "oracle_regret": 1.0, + "expert_regret": 0.8, + "action_mse_to_best": 0.3, + "restore_error": 2e-7, + }, + ] + + summary = _summarize_rows(rows) + + assert summary["num_groups"] == 2 + assert summary["policy_rollout_success_rate"] == 0.5 + assert summary["policy_rollout_progress"] == 0.5 + assert summary["oracle_success_rate"] == 1.0 + assert summary["expert_success_rate"] == 0.5 + assert summary["policy_oracle_regret"] == 0.6 + assert summary["restore_max_error"] == 2e-7 + + +def test_policy_rollout_loads_state_archive(tmp_path: Path) -> None: + archive = {"format": "dovla_maniskill_state_archive", "initial": {"g0": {"actors": {}}}} + (tmp_path / "state_archive.pkl").write_bytes(pickle.dumps(archive)) + + loaded = _load_state_archive(tmp_path) + + assert loaded["initial"]["g0"] == {"actors": {}} + + +def test_policy_rollout_requires_numeric_action_values() -> None: + record = CILRecord( + version=CIL_VERSION, + record_id="r0", + group_id="g0", + state_hash="h0", + task_id="PickCube-v1", + scene_id=None, + instruction="pick", + instruction_family={}, + observation_ref=None, + observation_inline={"features": [0.0]}, + action_chunk=ActionChunk( + representation="semantic", + horizon=1, + values=[{"command": "grasp", "object": "cube"}], + ), + next_observation_ref=None, + next_observation_inline={"features": [0.0]}, + structured_effect=StructuredEffect(), + reward=RewardInfo(progress=0.0, success=False, terminal_success=False), + regret=None, + rank_within_group=None, + candidate_type="expert", + failure=FailureInfo(type="unknown"), + ) + + try: + _numeric_action_values(record) + except ValueError as exc: + assert "numeric action chunks" in str(exc) + else: # pragma: no cover + raise AssertionError("symbolic action chunks cannot be used for ManiSkill rollout") diff --git a/tests/test_maniskill_render.py b/tests/test_maniskill_render.py new file mode 100644 index 0000000000000000000000000000000000000000..909d2b2a0772ad68ef2fbeba61d88efe25bc2bcb --- /dev/null +++ b/tests/test_maniskill_render.py @@ -0,0 +1,52 @@ +from __future__ import annotations + +import json + +import pytest + +from dovla_cil.generation.maniskill_render import ( + _load_records, + _rewrite_observation_refs, + _validate_state_archive, +) + + +def test_offline_renderer_rewrites_grouped_observation_refs(tmp_path) -> None: + shard = tmp_path / "shard_000000.jsonl" + records = [ + {"record_id": "r0", "group_id": "g0", "observation_ref": None}, + {"record_id": "r1", "group_id": "g0", "observation_ref": None}, + {"record_id": "r2", "group_id": "g1", "observation_ref": None}, + ] + shard.write_text( + "".join(json.dumps(record) + "\n" for record in records), + encoding="utf-8", + ) + + by_group, by_shard = _load_records([shard]) + _rewrite_observation_refs( + by_shard, + { + "r0": ("observations.h5#initial_rgb_jpeg/0", "next/0"), + "r1": ("observations.h5#initial_rgb_jpeg/0", "next/1"), + "r2": ("observations.h5#initial_rgb_jpeg/1", "next/2"), + }, + ) + + rewritten = [json.loads(line) for line in shard.read_text().splitlines()] + assert list(by_group) == ["g0", "g1"] + assert [len(group) for group in by_group.values()] == [2, 1] + assert rewritten[0]["observation_ref"] == rewritten[1]["observation_ref"] + assert rewritten[2]["observation_ref"].endswith("/1") + assert [record["next_observation_ref"] for record in rewritten] == [ + "next/0", + "next/1", + "next/2", + ] + + +def test_offline_renderer_requires_versioned_before_after_states() -> None: + _validate_state_archive({"version": 2, "initial": {}, "next": {}}) + + with pytest.raises(ValueError, match="version 2"): + _validate_state_archive({"group-id": {"actors": {}}}) diff --git a/tests/test_openvla_adapter.py b/tests/test_openvla_adapter.py new file mode 100644 index 0000000000000000000000000000000000000000..5559e5c1b6a70eb401587f2be84732454f293c31 --- /dev/null +++ b/tests/test_openvla_adapter.py @@ -0,0 +1,159 @@ +from __future__ import annotations + +from pathlib import Path +from types import SimpleNamespace + +import pytest + +torch = pytest.importorskip("torch") + +from dovla_cil.data.schema import ActionChunk +from dovla_cil.models.dovla import DoVLAConfig, DoVLAModel +from dovla_cil.models.openvla_adapter import ( + ExternalOpenVLAAdapter, + PretrainedCLIPBackbone, + ToyVLABackbone, + VLABackbone, +) + + +def _config() -> DoVLAConfig: + return DoVLAConfig( + obs_dim=10, + lang_dim=16, + action_dim=8, + hidden_dim=32, + action_horizon=3, + effect_dim=7, + intervention_dim=24, + ) + + +def test_toy_vla_backbone_works() -> None: + config = _config() + backbone = ToyVLABackbone(config) + observation = torch.randn(2, config.obs_dim) + instructions = ["pick the mug", "open the drawer"] + action = torch.randn(2, config.action_horizon, config.action_dim) + + context = backbone.encode_observation_language(observation, instructions) + action_z = backbone.encode_action(action) + policy = backbone.forward_policy(observation, instructions) + intervention = backbone.forward_intervention(observation, instructions, action) + decoded = backbone.decode_action(policy) + + assert isinstance(backbone, VLABackbone) + assert context.shape == (2, config.hidden_dim) + assert action_z.shape == (2, config.hidden_dim) + assert policy.shape == (2, config.action_horizon, config.action_dim) + assert intervention.shape == (2, config.intervention_dim) + assert isinstance(decoded, ActionChunk) + + +def test_external_openvla_adapter_requires_configuration() -> None: + with pytest.raises(NotImplementedError) as exc_info: + ExternalOpenVLAAdapter() + + assert "checkpoint_path" in str(exc_info.value) + + +def test_external_openvla_adapter_methods_are_placeholders() -> None: + adapter = ExternalOpenVLAAdapter(checkpoint_path=Path("missing-openvla-checkpoint")) + + with pytest.raises(NotImplementedError) as exc_info: + adapter.forward_policy(None, "pick the mug") + + assert "extension point only" in str(exc_info.value) + + +def test_dovla_model_can_use_injected_backbone() -> None: + config = _config() + backbone = ToyVLABackbone(config) + model = DoVLAModel(config, backbone=backbone) + observation = torch.randn(2, config.obs_dim) + instructions = ["pick the mug", "open the drawer"] + action = torch.randn(2, config.action_horizon, config.action_dim) + + policy = model.forward_policy(observation, instructions) + effect = model.forward_effect(observation, instructions, action) + reward = model.forward_reward(observation, instructions, action) + z = model.encode_intervention(observation, instructions, action) + + assert model.backbone is backbone + assert policy.shape == (2, config.action_horizon, config.action_dim) + assert effect["effect_vector"].shape == (2, config.effect_dim) + assert reward.shape == (2,) + assert z.shape == (2, config.intervention_dim) + + +class _FakeCLIP(torch.nn.Module): + def __init__(self, projection_dim: int = 12) -> None: + super().__init__() + self.anchor = torch.nn.Parameter(torch.zeros(())) + self.config = SimpleNamespace( + projection_dim=projection_dim, + vision_config=SimpleNamespace(image_size=8), + ) + + def get_image_features(self, *, pixel_values): + values = pixel_values.mean(dim=(1, 2, 3), keepdim=False).unsqueeze(1) + return values.repeat(1, self.config.projection_dim) + self.anchor + + def get_text_features(self, *, input_ids, **_kwargs): + values = input_ids.float().mean(dim=1, keepdim=True) + return values.repeat(1, self.config.projection_dim) + self.anchor + + +class _FakeProcessor: + tokenizer = None + + def __init__(self) -> None: + self.tokenizer = self + + def __call__(self, texts, **_kwargs): + return { + "input_ids": torch.tensor( + [[len(word) for word in text.split()][:4] for text in texts], + dtype=torch.long, + ) + } + + +def test_pretrained_clip_backbone_supports_raw_and_cached_features() -> None: + config = DoVLAConfig( + obs_dim=10, + lang_dim=16, + action_dim=8, + hidden_dim=32, + action_horizon=3, + effect_dim=7, + intervention_dim=24, + observation_mode="rgb", + backbone_type="native", + ) + backbone = PretrainedCLIPBackbone( + config, + clip_model=_FakeCLIP(), + processor=_FakeProcessor(), + ) + images = torch.randint(0, 255, (2, 12, 10, 3), dtype=torch.uint8) + instructions = ["pick red cube", "push blue cube"] + action = torch.randn(2, config.action_horizon, config.action_dim) + + cached = backbone.encode_pretrained_features(images, instructions) + raw_context = backbone.encode_observation_language(images, instructions) + cached_context = backbone.encode_observation_language(cached, instructions) + intervention = backbone.forward_intervention(cached, instructions, action) + + assert cached.shape == (2, 24) + assert torch.allclose(raw_context, cached_context) + assert raw_context.shape == (2, config.hidden_dim) + assert intervention.shape == (2, config.intervention_dim) + assert not any(parameter.requires_grad for parameter in backbone.clip_model.parameters()) + + +def test_clip_model_config_requires_rgb_and_model_path() -> None: + with pytest.raises(ValueError, match="observation_mode"): + DoVLAConfig(backbone_type="clip", backbone_model="local-clip") + with pytest.raises(ValueError, match="backbone_model"): + DoVLAConfig(observation_mode="rgb", backbone_type="clip") diff --git a/tests/test_paper_artifacts.py b/tests/test_paper_artifacts.py new file mode 100644 index 0000000000000000000000000000000000000000..37516a54164373e776268898cb61c35317d29787 --- /dev/null +++ b/tests/test_paper_artifacts.py @@ -0,0 +1,164 @@ +from __future__ import annotations + +import subprocess +import sys +from pathlib import Path + +from dovla_cil.utils.io import read_json, write_json + + +def test_make_paper_artifacts_with_fake_metrics(tmp_path: Path) -> None: + runs_dir = tmp_path / "runs" + out_dir = tmp_path / "paper_artifacts" + _write_fake_paper_runs(runs_dir) + + result = subprocess.run( + [ + sys.executable, + "scripts/make_paper_artifacts.py", + "--runs", + str(runs_dir), + "--out", + str(out_dir), + ], + check=True, + text=True, + capture_output=True, + ) + + assert "paper artifacts:" in result.stdout + for filename in ( + "main_scaling_table.csv", + "main_scaling_table.md", + "baseline_comparison_table.csv", + "baseline_comparison_table.md", + "ablation_table.csv", + "ablation_table.md", + "causalstress_per_category_table.csv", + "causalstress_per_category_table.md", + "result_summary.md", + "artifact_manifest.json", + ): + assert (out_dir / filename).exists() + + for filename in ( + "performance_vs_k.png", + "same_state_vs_cross_state_ranking.png", + "physical_outcome_vs_label_only.png", + "success_by_failure_category.png", + "regret_calibration.png", + ): + path = out_dir / "figures" / filename + assert path.exists() + assert path.stat().st_size > 0 + + summary = (out_dir / "result_summary.md").read_text(encoding="utf-8") + manifest = read_json(out_dir / "artifact_manifest.json") + assert "Best detected model" in summary + assert "Expected Claim Checks" in summary + assert manifest["num_scaling_rows"] == 2 + assert manifest["num_baseline_rows"] >= 4 + assert manifest["num_category_rows"] == 2 + + +def test_paper_artifacts_loads_measured_lattice_eval(tmp_path: Path) -> None: + from scripts.make_paper_artifacts import collect_metric_rows, load_result_payloads + + runs = tmp_path / "runs" + write_json( + { + "k": 8, + "selected_success_rate": 0.75, + "pairwise_ranking_accuracy": 0.8, + "top1_action_selection": 0.7, + "effect_prediction_mae": 0.2, + }, + runs / "scaling" / "k_8" / "seed_0" / "lattice_eval.json", + ) + + rows = collect_metric_rows(load_result_payloads(runs), runs) + + assert len(rows) == 1 + assert rows[0]["k"] == 8 + assert rows[0]["success_rate"] == 0.75 + + +def _write_fake_paper_runs(root: Path) -> None: + scaling_payloads = [ + { + "run_name": "k1", + "k": 1, + "num_states": 16, + "effective_total_records": 16, + "task_success_rate": 0.35, + "pairwise_ranking_accuracy": 0.50, + "top1_action_selection": 0.45, + "instruction_switch_accuracy": 0.30, + "effect_prediction_mae": 0.80, + "regret_calibration_error": 0.25, + }, + { + "run_name": "k4", + "k": 4, + "num_states": 4, + "effective_total_records": 16, + "task_success_rate": 0.65, + "pairwise_ranking_accuracy": 0.78, + "top1_action_selection": 0.70, + "instruction_switch_accuracy": 0.55, + "effect_prediction_mae": 0.50, + "regret_calibration_error": 0.12, + }, + ] + for payload in scaling_payloads: + write_json(payload, root / "scaling_toy" / f"k_{payload['k']:04d}" / "metrics.json") + + baselines = { + "expert_only_bc": (0.45, 0.55), + "cross_state_negatives": (0.40, 0.45), + "label_only_counterfactual": (0.42, 0.48), + "no_rank_regret": (0.50, 0.52), + "world_model_auxiliary": (0.54, 0.56), + } + for baseline, (success, ranking) in baselines.items(): + write_json( + { + "baseline": baseline, + "eval": { + "task_success_rate": success, + "pairwise_ranking_accuracy": ranking, + "top1_action_selection": success, + "instruction_switch_accuracy": success - 0.05, + "effect_prediction_mae": 1.0 - success, + "regret_calibration_error": 1.0 - ranking, + }, + }, + root / "baselines" / baseline / "metrics.json", + ) + + write_json( + { + "run_name": "causalstress_best", + "task_success_rate": 0.66, + "pairwise_ranking_accuracy": 0.79, + "per_category": { + "wrong_target_distractor": { + "success": 0.60, + "selected_success": 0.70, + "failure_rate": 0.40, + "instruction_switch": 0.75, + "top1": 0.80, + "pair_correct": 0.79, + }, + "near_miss_boundary": { + "success": 0.55, + "selected_success": 0.65, + "failure_rate": 0.45, + "instruction_switch": 0.70, + "top1": 0.76, + "pair_correct": 0.74, + }, + }, + }, + root / "dovla_toy" / "causalstress.json", + ) diff --git a/tests/test_retrieval.py b/tests/test_retrieval.py new file mode 100644 index 0000000000000000000000000000000000000000..5865bc22b7dc2050d4b2c6cd62a13cfbb6f9db23 --- /dev/null +++ b/tests/test_retrieval.py @@ -0,0 +1,137 @@ +from __future__ import annotations + +from pathlib import Path + +from dovla_cil.data.datasets import CILDataset +from dovla_cil.generation.pipeline import generate_cil_dataset +from dovla_cil.retrieval.embeddings import cosine_similarity, embed_observation_language +from dovla_cil.retrieval.eval import RetrievalEvalQuery, evaluate_retrieval_baselines +from dovla_cil.retrieval.index import CILRetrievalIndex +from dovla_cil.retrieval.prompting import build_retrieval_prompt +from dovla_cil.retrieval.retriever import CriticGatedRetriever, RetrievalConditionedPolicyWrapper +from dovla_cil.tasks.library import built_in_toy_tasks +from dovla_cil.transfercritic.schema import TransferContext + + +def _make_dataset(tmp_path: Path) -> CILDataset: + generate_cil_dataset( + backend="toy", + tasks=built_in_toy_tasks()[:2], + out_dir=tmp_path, + num_states_per_task=1, + k=4, + seed=21, + shard_size=8, + inline_observations=True, + ) + return CILDataset(tmp_path) + + +def test_embedding_index_over_tiny_cil_dataset(tmp_path: Path) -> None: + dataset = _make_dataset(tmp_path) + index = CILRetrievalIndex.from_dataset(tmp_path, dim=32) + record = dataset[0] + query = embed_observation_language(record.observation_inline, record.instruction, dim=32) + + hits = index.query(query, top_k=3) + + assert len(index) == len(dataset) + assert len(query) == 32 + assert len(hits) == 3 + assert hits[0].similarity >= hits[-1].similarity + assert cosine_similarity(query, query) > 0.999 + + +def test_retriever_modes_and_same_state_filter(tmp_path: Path) -> None: + dataset = _make_dataset(tmp_path) + index = CILRetrievalIndex.from_dataset(tmp_path, dim=32) + first = dataset[0] + retriever = CriticGatedRetriever(index, dim=32) + + same_state = retriever.retrieve( + first.observation_inline, + first.instruction, + k=3, + mode="nearest_neighbor", + same_state_group_id=first.group_id, + ) + success_only = retriever.retrieve(first.observation_inline, first.instruction, k=3, mode="success_only") + contrastive = retriever.retrieve( + first.observation_inline, + first.instruction, + k=3, + mode="success_failure_contrastive", + ) + + assert same_state.examples + assert all(example.item.group_id == first.group_id for example in same_state.examples) + assert all(example.item.success for example in success_only.examples) + assert contrastive.examples + assert any(example.role == "positive_successful" for example in contrastive.examples) + + +def test_critic_gated_retrieval_uses_optional_critic(tmp_path: Path) -> None: + dataset = _make_dataset(tmp_path) + index = CILRetrievalIndex.from_dataset(tmp_path, dim=32) + first = dataset[0] + + class MockCritic: + def score_atom(self, atom, selected_atoms, context): + del selected_atoms, context + return 10.0 if atom.reward_summary.get("success", 0.0) else 0.0 + + retriever = CriticGatedRetriever( + index, + critic=MockCritic(), + transfer_context=TransferContext(benchmark_name="CausalStress"), + dim=32, + ) + result = retriever.retrieve(first.observation_inline, first.instruction, k=3, mode="critic_gated") + + assert result.examples + assert result.examples[0].gate_score >= result.examples[-1].gate_score + + +def test_retrieval_conditioned_policy_wrapper(tmp_path: Path) -> None: + dataset = _make_dataset(tmp_path) + index = CILRetrievalIndex.from_dataset(tmp_path, dim=32) + first = dataset[0] + retriever = CriticGatedRetriever(index, dim=32) + + class DummyPolicy: + def forward_policy(self, observation, instruction, retrieved_examples=None): + del observation, instruction + return {"retrieved": len(retrieved_examples or [])} + + wrapper = RetrievalConditionedPolicyWrapper(DummyPolicy(), retriever, k=2) + output = wrapper.policy(first.observation_inline, first.instruction) + + assert output["retrieved"] == 2 + assert len(wrapper.last_retrieved_examples) == 2 + prompt = build_retrieval_prompt(first.instruction, wrapper.last_retrieved_examples) + assert "Retrieved exemplars" in prompt + + +def test_retrieval_eval_baselines(tmp_path: Path) -> None: + dataset = _make_dataset(tmp_path) + index = CILRetrievalIndex.from_dataset(tmp_path, dim=32) + retriever = CriticGatedRetriever(index, dim=32) + queries = [ + RetrievalEvalQuery( + observation=record.observation_inline or {}, + instruction=record.instruction, + group_id=record.group_id, + ) + for record in dataset.records[:2] + ] + + report = evaluate_retrieval_baselines(retriever, queries, k=3) + + assert set(report) == { + "no_retrieval", + "nearest_neighbor", + "success_only", + "success_failure_contrastive", + "critic_gated", + } + assert report["nearest_neighbor"]["retrieval_coverage"] == 1.0 diff --git a/tests/test_scaling.py b/tests/test_scaling.py new file mode 100644 index 0000000000000000000000000000000000000000..4504f62a0c86004623a03a66b0a892ab93259b90 --- /dev/null +++ b/tests/test_scaling.py @@ -0,0 +1,54 @@ +from __future__ import annotations + +from pathlib import Path + +from dovla_cil.experiments.scaling import ( + ScalingExperiment, + parse_k_values, + read_scaling_csv, + run_scaling_experiment, +) +from dovla_cil.utils.io import read_json + + +def test_parse_k_values() -> None: + assert parse_k_values("1,2,4") == (1, 2, 4) + + +def test_tiny_scaling_run_writes_csv_and_plots(tmp_path: Path) -> None: + summary = run_scaling_experiment( + ScalingExperiment( + backend="toy", + tasks="builtins", + output_dir=tmp_path, + total_records=4, + k_values=(1, 2), + epochs=1, + seed=0, + shard_size=8, + batch_groups=2, + records_per_group=2, + hidden_dim=32, + eval_num_tasks=2, + device="auto", + ) + ) + + rows = read_scaling_csv(tmp_path / "scaling_results.csv") + regression = read_json(tmp_path / "scaling_regression.json") + + assert len(rows) == 2 + assert [row["k"] for row in rows] == ["1", "2"] + assert Path(summary["aggregate_csv"]).exists() + assert "ranking_acc" in regression + for filename in ( + "success_rate_vs_k.png", + "ranking_acc_vs_k.png", + "instruction_switch_acc_vs_k.png", + "effect_mae_vs_k.png", + "regret_ece_vs_k.png", + ): + assert (tmp_path / filename).exists() + assert (tmp_path / filename).stat().st_size > 0 + assert (tmp_path / "k_0001" / "metrics.json").exists() + assert (tmp_path / "k_0002" / "metrics.json").exists() diff --git a/tests/test_slurm_templates.py b/tests/test_slurm_templates.py new file mode 100644 index 0000000000000000000000000000000000000000..4dacb6c3f7dd98d1cec0262a233af8204b58d2e1 --- /dev/null +++ b/tests/test_slurm_templates.py @@ -0,0 +1,110 @@ +from __future__ import annotations + +from pathlib import Path + + +def test_visual_lattice_eval_array_skips_indices_beyond_visual_seed_count() -> None: + script = Path("scripts/slurm/eval_lattice_array.sbatch").read_text(encoding="utf-8") + + assert '#SBATCH --array=0-5%6' in script + assert 'VISUAL_SEEDS="${VISUAL_SEEDS:-3}"' in script + assert "TASK_INDEX >= VISUAL_SEEDS" in script + assert 'exit 0' in script + + +def test_lattice_eval_array_supports_field_only_transfer_mode() -> None: + script = Path("scripts/slurm/eval_lattice_array.sbatch").read_text(encoding="utf-8") + + assert 'MODE" == "field_only"' in script + assert 'OBJECTIVE="${OBJECTIVE:-lattice_field}"' in script + assert 'RUN_DIR="$RUN_ROOT/$OBJECTIVE/seed_$SEED"' in script + assert 'ALL_GROUPS:-0' in script + assert 'EVAL_EXTRA_ARGS+=(--all-groups)' in script + + +def test_collection_train_array_supports_field_only_transfer_mode() -> None: + script = Path("scripts/slurm/train_maniskill_collection_array.sbatch").read_text( + encoding="utf-8" + ) + + assert 'OBJECTIVE_MODE="${OBJECTIVE_MODE:-paired}"' in script + assert 'EPOCHS="${EPOCHS:-50}"' in script + assert 'BATCH_GROUPS="${BATCH_GROUPS:-32}"' in script + assert 'HIDDEN_DIM="${HIDDEN_DIM:-256}"' in script + assert 'OBJECTIVE_MODE" == "field_only"' in script + assert 'EXPECTED_GROUPS:-3500' in script + assert 'EXPECTED_RECORDS:-56000' in script + + +def test_smolvla_download_template_is_pinned_and_secret_free() -> None: + script = Path("scripts/slurm/download_smolvla_checkpoint.sbatch").read_text( + encoding="utf-8" + ) + + assert 'REPO_ID="${REPO_ID:-lerobot/smolvla_base}"' in script + assert "c83c3163b8ca9b7e67c509fffd9121e66cb96205" in script + assert 'DRY_RUN="${DRY_RUN:-0}"' in script + assert "--dry-run" in script + assert "dovla_download_manifest.json" in script + assert 'CA_BUNDLE="${CA_BUNDLE:-$SCRATCH_ROOT/ca-bundle.crt}"' in script + assert "SSL_CERT_FILE=$CA_BUNDLE" in script + assert "REQUESTS_CA_BUNDLE=$CA_BUNDLE" in script + assert 'HF_HUB_ETAG_TIMEOUT="${HF_HUB_ETAG_TIMEOUT:-20}"' in script + assert 'HF_HUB_DOWNLOAD_TIMEOUT="${HF_HUB_DOWNLOAD_TIMEOUT:-30}"' in script + assert "SmolVLA download preflight" in script + assert "Network is unreachable" in script + assert "apptainer" in script + assert "HF_TOKEN" not in script + assert "--token" not in script + + +def test_smolvla_smoke_template_is_offline_and_uses_isolated_env() -> None: + script = Path("scripts/slurm/smoke_smolvla_checkpoint.sbatch").read_text(encoding="utf-8") + + assert 'PYTHON="${PYTHON:-$SCRATCH_ROOT/envs/smolvla/bin/python}"' in script + assert "HF_HUB_OFFLINE=1" in script + assert "TRANSFORMERS_OFFLINE=1" in script + assert "scripts/smoke_smolvla_checkpoint.py" in script + assert "7b375e1b73b11138ff12fe22c8f2822d8fe03467" in script + assert '--vlm-metadata "$VLM_METADATA"' in script + assert 'CONTAINER_OUT="/workspace/${OUT#"$PROJECT_DIR/"}"' in script + assert "--device cuda" in script + + +def test_smolvla_env_installer_is_offline_and_pinned() -> None: + script = Path("scripts/slurm/install_smolvla_env.sbatch").read_text(encoding="utf-8") + + assert "lerobot-0.4.3-py3-none-any.whl" in script + assert "transformers==4.57.6+computecanada" in script + assert "huggingface-hub==0.35.3+computecanada" in script + assert "--no-index" in script + assert "--no-deps" in script + assert "/cvmfs/soft.computecanada.ca/custom/python/wheelhouse" in script + assert "pip install" in script + + +def test_smolvla_cil_baseline_template_is_offline_and_claim_scoped() -> None: + script = Path("scripts/slurm/run_smolvla_cil_baseline.sbatch").read_text( + encoding="utf-8" + ) + config = Path("configs/external/smolvla_cil_smoke.json").read_text(encoding="utf-8") + full_config = Path("configs/external/smolvla_cil_full.json").read_text(encoding="utf-8") + aligned_config = Path("configs/external/smolvla_cil_aligned.json").read_text( + encoding="utf-8" + ) + + assert "HF_HUB_OFFLINE=1" in script + assert "TRANSFORMERS_OFFLINE=1" in script + assert "gpu:nvidia_h100_80gb_hbm3_3g.40gb:1" in script + assert "dovla_cil.eval.smolvla_cil_baseline:run_smolvla_cil_baseline" in script + assert "smolvla_cil_smoke.json" in script + assert 'CONTAINER_ADAPTER_CONFIG="/workspace/${ADAPTER_CONFIG#"$PROJECT_DIR/"}"' in script + assert '"steps": 2' in config + assert '"max_eval_groups": 12' in config + assert '"steps": 1000' in full_config + assert '"batch_size": 4' in full_config + assert "external_vla_export_full_balanced" in full_config + assert '"split_mode": "dataset_group_shuffle"' in aligned_config + assert '"max_eval_groups": 700' in aligned_config + assert "API_KEY" not in script + config + assert "HF_TOKEN" not in script + config diff --git a/tests/test_smoke_full_pipeline.py b/tests/test_smoke_full_pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..247fe78517ae8b3ddab1272dd9f4c6e025a380c9 --- /dev/null +++ b/tests/test_smoke_full_pipeline.py @@ -0,0 +1,56 @@ +from __future__ import annotations + +import subprocess +import sys +from pathlib import Path + +from dovla_cil.utils.io import read_json + + +def test_reduced_full_smoke_pipeline_runs(tmp_path: Path) -> None: + out_dir = tmp_path / "smoke_full" + + result = subprocess.run( + [ + sys.executable, + "scripts/smoke_full_pipeline.py", + "--out", + str(out_dir), + "--num-tasks", + "1", + "--states-per-task", + "1", + "--k", + "2", + "--eval-num-tasks", + "1", + "--epochs", + "1", + "--batch-groups", + "1", + "--records-per-group", + "2", + "--hidden-dim", + "32", + "--device", + "cpu", + ], + check=True, + text=True, + capture_output=True, + ) + + assert "Final paths" in result.stdout + assert (out_dir / "tasks.jsonl").exists() + assert (out_dir / "cil_toy" / "manifest.json").exists() + assert (out_dir / "inspect.txt").exists() + assert (out_dir / "train" / "best.pt").exists() + assert (out_dir / "causalstress" / "metrics.json").exists() + assert (out_dir / "dataset_report" / "summary.json").exists() + assert (out_dir / "eval_report" / "report.md").exists() + + dataset_summary = read_json(out_dir / "dataset_report" / "summary.json") + eval_summary = read_json(out_dir / "eval_report" / "summary.json") + assert dataset_summary["num_records"] == 2 + assert dataset_summary["num_groups"] == 1 + assert eval_summary["num_runs"] == 1 diff --git a/tests/test_smolvla_checkpoint_smoke.py b/tests/test_smolvla_checkpoint_smoke.py new file mode 100644 index 0000000000000000000000000000000000000000..2868ba5106926dc94281de20eabac0bbe1e902ea --- /dev/null +++ b/tests/test_smolvla_checkpoint_smoke.py @@ -0,0 +1,62 @@ +from __future__ import annotations + +import json +import subprocess +import sys +from pathlib import Path + +import pytest + +from scripts.smoke_smolvla_checkpoint import smoke_checkpoint + + +def _checkpoint(tmp_path: Path) -> Path: + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + (checkpoint / "config.json").write_text("{}", encoding="utf-8") + (checkpoint / "model.safetensors").write_bytes(b"weights") + return checkpoint + + +def test_metadata_only_smoke_does_not_require_lerobot(tmp_path: Path) -> None: + result = smoke_checkpoint(_checkpoint(tmp_path), metadata_only=True) + + assert result["schema_version"] == "smolvla-checkpoint-smoke/v0" + assert result["metadata_only"] is True + assert "lerobot" in result["package_versions"] + + +def test_smoke_rejects_incomplete_checkpoint(tmp_path: Path) -> None: + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + (checkpoint / "config.json").write_text("{}", encoding="utf-8") + + with pytest.raises(FileNotFoundError, match="model.safetensors"): + smoke_checkpoint(checkpoint, metadata_only=True) + + +def test_weight_smoke_requires_local_vlm_metadata(tmp_path: Path) -> None: + with pytest.raises(FileNotFoundError, match="--vlm-metadata"): + smoke_checkpoint(_checkpoint(tmp_path)) + + +def test_metadata_only_cli_writes_manifest(tmp_path: Path) -> None: + checkpoint = _checkpoint(tmp_path) + out = tmp_path / "smoke.json" + + result = subprocess.run( + [ + sys.executable, + "scripts/smoke_smolvla_checkpoint.py", + "--checkpoint", + str(checkpoint), + "--out", + str(out), + "--metadata-only", + ], + capture_output=True, + text=True, + ) + + assert result.returncode == 0 + assert json.loads(out.read_text(encoding="utf-8"))["metadata_only"] is True diff --git a/tests/test_smolvla_cil_baseline.py b/tests/test_smolvla_cil_baseline.py new file mode 100644 index 0000000000000000000000000000000000000000..0b6efbc0fa8aca5cec574aca4e6e380f5b5842e4 --- /dev/null +++ b/tests/test_smolvla_cil_baseline.py @@ -0,0 +1,191 @@ +from __future__ import annotations + +import random +from dataclasses import replace + +import numpy as np + +from dovla_cil.data.schema import ( + CIL_VERSION, + ActionChunk, + CILRecord, + RewardInfo, + StructuredEffect, + make_record_id, +) +from dovla_cil.eval.smolvla_cil_baseline import ( + candidate_selection_metrics, + dataset_group_row_split, + fit_action_normalizer, + fit_state_projector, + select_measured_candidate, + stratified_row_split, +) + + +def _record(record_id: str, action: list[list[float]], reward: float) -> CILRecord: + return CILRecord( + version=CIL_VERSION, + record_id=record_id, + group_id="group", + state_hash="state", + task_id="task", + scene_id="scene", + instruction="move the cube", + instruction_family={}, + observation_ref=None, + observation_inline={"features": [0.0]}, + action_chunk=ActionChunk( + action_id=make_record_id("group", record_id, seed=0), + representation="numeric", + horizon=len(action), + values=action, + ), + next_observation_ref=None, + next_observation_inline={"features": [0.0]}, + structured_effect=StructuredEffect(), + reward=RewardInfo( + progress=reward, + success=reward >= 1.0, + terminal_success=reward >= 1.0, + ), + regret=None, + rank_within_group=None, + candidate_type="test", + failure=None, + ) + + +def test_state_projector_and_action_normalizer_are_train_only_transforms() -> None: + states = np.arange(40, dtype=np.float32).reshape(10, 4) + projector = fit_state_projector(states, output_dim=3) + actions = np.arange(60, dtype=np.float32).reshape(10, 2, 3) + normalizer = fit_action_normalizer(actions) + + assert projector.transform(states[:2]).shape == (2, 3) + normalized = normalizer.normalize(actions) + assert np.allclose(normalized.reshape(-1, 3).mean(axis=0), 0.0, atol=1e-6) + assert np.allclose(normalizer.denormalize(normalized), actions) + + +def test_state_projector_handles_variable_task_state_dimensions() -> None: + states = [ + np.arange(4, dtype=np.float32), + np.arange(2, dtype=np.float32), + np.arange(3, dtype=np.float32), + np.arange(4, dtype=np.float32) + 1, + ] + + projector = fit_state_projector(states, output_dim=2) + + assert projector.raw_input_dim == 4 + assert projector.mean.shape == (8,) + assert projector.transform(states[1]).shape == (2,) + assert projector.to_dict()["state_vectorization"] == ( + "right_zero_pad_plus_validity_mask" + ) + + +def test_stratified_row_split_is_deterministic_and_task_covered() -> None: + rows = [ + {"cil": {"task_id": task, "group_id": f"{task}-{index}"}} + for task in ("pick", "push") + for index in range(5) + ] + + first_train, first_val = stratified_row_split(rows, val_fraction=0.2, seed=9) + second_train, second_val = stratified_row_split(list(reversed(rows)), val_fraction=0.2, seed=9) + + assert {row["cil"]["task_id"] for row in first_val} == {"pick", "push"} + assert [row["cil"]["group_id"] for row in first_train] == [ + row["cil"]["group_id"] for row in second_train + ] + assert [row["cil"]["group_id"] for row in first_val] == [ + row["cil"]["group_id"] for row in second_val + ] + + +def test_dataset_group_row_split_matches_global_group_order() -> None: + group_ids = [f"group-{index}" for index in range(10)] + rows = [ + {"cil": {"task_id": "task", "group_id": group_id}} + for group_id in reversed(group_ids) + ] + + train, validation = dataset_group_row_split( + rows, + group_ids, + val_fraction=0.2, + seed=3, + ) + + expected = list(group_ids) + random.Random(3).shuffle(expected) + assert [row["cil"]["group_id"] for row in validation] == expected[:2] + assert [row["cil"]["group_id"] for row in train] == expected[2:] + + +def test_candidate_selection_uses_measured_same_state_outcome() -> None: + records = [ + _record("low", [[0.0, 0.0], [0.0, 0.0]], 0.0), + _record("best", [[1.0, 1.0], [1.0, 1.0]], 1.0), + _record("middle", [[0.5, 0.5], [0.5, 0.5]], 0.5), + ] + actions = np.stack( + [np.asarray(record.action_chunk.values, dtype=np.float32) for record in records] + ) + normalizer = fit_action_normalizer(actions) + best = max(records, key=lambda record: record.reward.score) + prediction = np.asarray(best.action_chunk.values, dtype=np.float32) + group_id = records[0].group_id + + metrics = candidate_selection_metrics( + {group_id: prediction}, + {group_id: records}, + normalizer, + ) + + assert metrics["num_eval_groups"] == 1 + assert metrics["top1_action_selection"] == 1.0 + assert metrics["selected_reward_mean"] == best.reward.score + assert metrics["normalized_action_mse_to_selected"] == 0.0 + + +def test_candidate_selection_breaks_distance_ties_by_record_id() -> None: + first = _record("zzz-record", [[0.25, -0.25]], 0.2) + duplicate = replace( + first, + record_id="aaa-record", + action_chunk=ActionChunk( + action_id="duplicate", + representation=first.action_chunk.representation, + horizon=first.action_chunk.horizon, + values=first.action_chunk.values, + ), + ) + records = [first, duplicate] + action = np.asarray(first.action_chunk.values, dtype=np.float32) + normalizer = fit_action_normalizer(np.stack([action, action])) + + selected, distance = select_measured_candidate(action, records, normalizer) + + assert selected.record_id == "aaa-record" + assert distance == 0.0 + + +def test_candidate_selection_pads_mixed_action_dimensions() -> None: + short = _record("short", [[0.0, 0.0]], 0.2) + wide = _record("wide", [[1.0, 1.0, 1.0]], 1.0) + canonical_actions = np.asarray( + [[[0.0, 0.0, 0.0]], [[1.0, 1.0, 1.0]]], + dtype=np.float32, + ) + normalizer = fit_action_normalizer(canonical_actions) + + selected, _ = select_measured_candidate( + canonical_actions[0], + [short, wide], + normalizer, + ) + + assert selected.record_id == "short" diff --git a/tests/test_smolvla_runtime.py b/tests/test_smolvla_runtime.py new file mode 100644 index 0000000000000000000000000000000000000000..45edd5c551cf7bb976f789ff9226f35f57d42f20 --- /dev/null +++ b/tests/test_smolvla_runtime.py @@ -0,0 +1,88 @@ +from __future__ import annotations + +import sys +import types +from collections import deque +from pathlib import Path + +from dovla_cil.eval.smolvla_runtime import ( + ensure_lazy_policies_namespace, + ensure_policy_utils_shim, + ensure_train_config_type_shim, +) + + +def test_lazy_smolvla_namespace_points_to_installed_policy_directory( + tmp_path: Path, + monkeypatch, +) -> None: + package_dir = tmp_path / "lerobot" + policies_dir = package_dir / "policies" + policies_dir.mkdir(parents=True) + module = types.ModuleType("lerobot") + module.__file__ = str(package_dir / "__init__.py") + monkeypatch.delitem(sys.modules, "lerobot.policies", raising=False) + + namespace = ensure_lazy_policies_namespace(module) + + assert namespace.__path__ == [str(policies_dir)] + assert namespace.__spec__ is not None + assert namespace.__spec__.submodule_search_locations == [str(policies_dir)] + assert sys.modules["lerobot.policies"] is namespace + + +def test_lazy_smolvla_namespace_is_idempotent(tmp_path: Path, monkeypatch) -> None: + package_dir = tmp_path / "lerobot" + (package_dir / "policies").mkdir(parents=True) + module = types.ModuleType("lerobot") + module.__file__ = str(package_dir / "__init__.py") + monkeypatch.delitem(sys.modules, "lerobot.policies", raising=False) + + first = ensure_lazy_policies_namespace(module) + second = ensure_lazy_policies_namespace(module) + + assert first is second + + +def test_train_config_shim_exposes_only_required_type(monkeypatch) -> None: + configs = types.ModuleType("lerobot.configs") + configs.__path__ = [] + monkeypatch.setitem(sys.modules, "lerobot.configs", configs) + monkeypatch.delitem(sys.modules, "lerobot.configs.train", raising=False) + + module = ensure_train_config_type_shim() + + assert module.__dovla_type_only_shim__ is True + assert module.TrainPipelineConfig.__module__ == "lerobot.configs.train" + assert configs.train is module + + +def test_train_config_shim_preserves_real_import(monkeypatch) -> None: + real_module = types.ModuleType("lerobot.configs.train") + monkeypatch.setitem(sys.modules, "lerobot.configs.train", real_module) + + assert ensure_train_config_type_shim() is real_module + + +def test_policy_utils_shim_populates_and_advances_queues(monkeypatch) -> None: + monkeypatch.delitem(sys.modules, "lerobot.policies.utils", raising=False) + module = ensure_policy_utils_shim() + queues = {"observation": deque(maxlen=2), "action": deque(maxlen=2)} + + result = module.populate_queues( + queues, + {"observation": "first", "action": "excluded"}, + exclude_keys=["action"], + ) + module.populate_queues(queues, {"observation": "second"}) + + assert result is queues + assert list(queues["observation"]) == ["first", "second"] + assert list(queues["action"]) == [] + + +def test_policy_utils_shim_preserves_real_import(monkeypatch) -> None: + real_module = types.ModuleType("lerobot.policies.utils") + monkeypatch.setitem(sys.modules, "lerobot.policies.utils", real_module) + + assert ensure_policy_utils_shim() is real_module diff --git a/tests/test_task_schema.py b/tests/test_task_schema.py new file mode 100644 index 0000000000000000000000000000000000000000..7b20cbdb0fb42242dbbd58996404df8ba581cd38 --- /dev/null +++ b/tests/test_task_schema.py @@ -0,0 +1,70 @@ +from __future__ import annotations + +import pytest + +from dovla_cil.tasks.library import ( + HARD_CAUSALSTRESS_CATEGORIES, + ToyTaskLibrary, + built_in_causalstress_tasks, + built_in_toy_tasks, +) +from dovla_cil.tasks.predicates import evaluate_predicate, evaluate_task_success +from dovla_cil.tasks.schema import RelationSpec, TaskSpec +from dovla_cil.tasks.validators import validate_task + + +def test_valid_builtin_tasks_pass() -> None: + tasks = built_in_toy_tasks() + assert len(tasks) == 10 + for task in tasks: + validate_task(task) + assert TaskSpec.from_dict(task.to_dict()) == task + + +def test_toy_task_library_cycles_tasks() -> None: + library = ToyTaskLibrary() + assert library.get(0).task_id == "toy_pick_red_mug" + assert library.get(10).task_id == "toy_pick_red_mug" + assert len(library.list()) == 10 + + +def test_hard_causalstress_tasks_pass_validation() -> None: + tasks = built_in_causalstress_tasks() + assert len(tasks) == len(HARD_CAUSALSTRESS_CATEGORIES) + assert {task.family for task in tasks} == set(HARD_CAUSALSTRESS_CATEGORIES) + for task in tasks: + validate_task(task) + + +def test_invalid_predicate_object_fails() -> None: + task = built_in_toy_tasks()[0].model_copy( + update={"success_predicates": [RelationSpec(name="grasped", args=["missing_mug"])]} + ) + with pytest.raises(ValueError, match="unknown object"): + validate_task(task) + + +def test_predicate_evaluator_inside_left_of_opened_lifted() -> None: + state = { + "objects": { + "red_mug": {"inside": "blue_bowl", "position": [0.0, 0.0, 0.2]}, + "blue_bowl": {"position": [1.0, 0.0, 0.0]}, + "drawer": {"opened": True}, + "can": {"position": [0.0, 0.0, 0.3]}, + } + } + assert evaluate_predicate(RelationSpec(name="inside", args=["red_mug", "blue_bowl"]), state) + assert evaluate_predicate(RelationSpec(name="left_of", args=["red_mug", "blue_bowl"]), state) + assert evaluate_predicate(RelationSpec(name="opened", args=["drawer"]), state) + assert evaluate_predicate(RelationSpec(name="lifted", args=["can"]), state) + + +def test_evaluate_task_success() -> None: + task = ToyTaskLibrary().get_by_id("toy_put_red_mug_in_blue_bowl") + state = { + "objects": { + "red_mug": {"inside": "blue_bowl", "position": [0.0, 0.0, 0.0]}, + "blue_bowl": {"position": [0.0, 0.0, 0.0]}, + } + } + assert evaluate_task_success(task, state) diff --git a/tests/test_toy_sim.py b/tests/test_toy_sim.py new file mode 100644 index 0000000000000000000000000000000000000000..aedbd475b2e2438d7cf8a1df73862446488dbe8a --- /dev/null +++ b/tests/test_toy_sim.py @@ -0,0 +1,82 @@ +from __future__ import annotations + +from dovla_cil.sim.base import ActionChunk, SimState +from dovla_cil.sim.registry import get_simulator_backend +from dovla_cil.sim.toy_backend import ToyBackend +from dovla_cil.tasks.library import ToyTaskLibrary +from dovla_cil.tasks.predicates import evaluate_task_success + + +def test_toy_backend_can_reset_task() -> None: + task = ToyTaskLibrary().get_by_id("toy_put_red_mug_in_blue_bowl") + sim = ToyBackend() + sim.seed(123) + state = sim.reset_task(task) + observation = sim.render_observation() + + assert isinstance(state, SimState) + assert state.task_id == task.task_id + assert "red_mug" in observation["objects"] + assert observation["robot"]["gripper"] == "open" + + +def test_serialize_restore_gives_identical_symbolic_state() -> None: + task = ToyTaskLibrary().get_by_id("toy_put_red_mug_in_blue_bowl") + sim = ToyBackend() + sim.seed(123) + sim.reset_task(task) + state_blob = sim.serialize_state() + expected = sim.get_symbolic_state() + + sim.execute_action_chunk(ActionChunk.single("push", object="red_mug", dx=0.2, dy=0.0)) + sim.restore_state(state_blob) + + assert sim.get_symbolic_state() == expected + + +def test_same_action_from_same_restored_state_is_deterministic() -> None: + task = ToyTaskLibrary().get_by_id("toy_push_cube_to_target_zone") + sim = ToyBackend() + sim.seed(17) + sim.reset_task(task) + state_blob = sim.serialize_state() + action = ActionChunk.single("push", object="cube", dx=0.1, dy=0.2) + + first = sim.execute_action_chunk(action) + sim.restore_state(state_blob) + second = sim.execute_action_chunk(action) + + assert first.after_state == second.after_state + assert first.observation == second.observation + assert first.contacts == second.contacts + assert first.reward == second.reward + + +def test_predicate_success_after_toy_actions() -> None: + task = ToyTaskLibrary().get_by_id("toy_put_red_mug_in_blue_bowl") + sim = ToyBackend() + sim.seed(4) + sim.reset_task(task) + + result = sim.execute_action_chunk( + ActionChunk( + [ + {"command": "move_to", "object": "red_mug"}, + {"command": "grasp", "object": "red_mug"}, + { + "command": "place_at", + "object": "red_mug", + "container": "blue_bowl", + "relation": "inside", + }, + ] + ) + ) + + assert evaluate_task_success(task, sim.get_symbolic_state()) + assert result.info["success"] is True + assert result.reward == 1.0 + + +def test_registry_toy_backend_works() -> None: + assert isinstance(get_simulator_backend("toy"), ToyBackend) diff --git a/tests/test_trainer.py b/tests/test_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..9326ac31ae62e45e18b524f4bc13af90231b87eb --- /dev/null +++ b/tests/test_trainer.py @@ -0,0 +1,106 @@ +from __future__ import annotations + +from pathlib import Path +from types import SimpleNamespace + +from dovla_cil.data.schema import RewardInfo +from dovla_cil.generation.pipeline import generate_cil_dataset +from dovla_cil.tasks.library import built_in_toy_tasks +from dovla_cil.training.trainer import ( + DoVLATrainer, + TrainerConfig, + _cross_state_pair_indices, + _reward_utility_values, +) +from dovla_cil.utils.io import read_json + + +def test_trainer_runs_one_epoch_and_writes_checkpoints(tmp_path: Path) -> None: + dataset_dir = tmp_path / "cil" + run_dir = tmp_path / "run" + generate_cil_dataset( + backend="toy", + tasks=built_in_toy_tasks()[:3], + out_dir=dataset_dir, + num_states_per_task=2, + k=4, + seed=5, + shard_size=8, + inline_observations=True, + ) + + result = DoVLATrainer( + TrainerConfig( + dataset_dir=dataset_dir, + output_dir=run_dir, + epochs=1, + batch_groups=2, + records_per_group=4, + hidden_dim=64, + learning_rate=1e-3, + seed=5, + device="cpu", + ) + ).train() + + assert (run_dir / "latest.pt").exists() + assert (run_dir / "best.pt").exists() + assert "rank_acc" in result["history"][0]["val"] + metrics = read_json(run_dir / "metrics.json") + assert "rank_acc" in metrics["history"][0]["val"] + + +def test_field_utility_includes_terminal_success_bonus() -> None: + records = [ + SimpleNamespace( + reward=RewardInfo( + progress=0.4, + success=False, + terminal_success=False, + ) + ), + SimpleNamespace( + reward=RewardInfo( + progress=0.4, + success=True, + terminal_success=True, + ) + ), + ] + + assert _reward_utility_values(records) == [0.4, 1.4] + + +def test_cross_state_pairs_preserve_task_and_reward_order() -> None: + records = [ + SimpleNamespace( + task_id="pick", + group_id=f"g{group}", + reward=SimpleNamespace(score=reward), + ) + for group, reward in ((0, 0.1), (0, 0.9), (1, 0.2), (1, 0.8), (2, 0.4)) + ] + + pairs = _cross_state_pair_indices(records, pair_count=12, seed=7) + + assert len(pairs) == 12 + for better, worse in pairs: + assert records[better].task_id == records[worse].task_id + assert records[better].group_id != records[worse].group_id + assert records[better].reward.score > records[worse].reward.score + + +def test_cross_state_scope_rejects_lattice_field_objective(tmp_path: Path) -> None: + try: + TrainerConfig( + dataset_dir=tmp_path, + output_dir=tmp_path / "out", + objective="lattice_field", + pair_scope="cross_state", + ) + except ValueError as exc: + assert "legacy objective" in str(exc) + else: # pragma: no cover - protects baseline semantics + raise AssertionError( + "cross-state pairs cannot silently leave same-state field edges active" + ) diff --git a/tests/test_training_dataset.py b/tests/test_training_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..333a74fd8bf799bc5ffc5998395d81f61a0ba7a9 --- /dev/null +++ b/tests/test_training_dataset.py @@ -0,0 +1,107 @@ +from __future__ import annotations + +from pathlib import Path + +from dovla_cil.data.datasets import CILDataset, write_cil_collection +from dovla_cil.data.group_sampler import GroupAwareBatchSampler +from dovla_cil.generation.pipeline import generate_cil_dataset +from dovla_cil.tasks.library import built_in_toy_tasks +from dovla_cil.training.collate import collate_cil_records + + +def _make_toy_dataset(tmp_path: Path) -> CILDataset: + generate_cil_dataset( + backend="toy", + tasks=built_in_toy_tasks()[:3], + out_dir=tmp_path, + num_states_per_task=1, + k=4, + seed=13, + shard_size=8, + inline_observations=True, + ) + return CILDataset(tmp_path) + + +def test_cil_dataset_loads_generated_toy_cil(tmp_path: Path) -> None: + dataset = _make_toy_dataset(tmp_path) + assert len(dataset) == 12 + first = dataset[0] + assert dataset.get_group(first.group_id)[0] == first + assert len(list(dataset.iter_groups())) == 3 + + +def test_full_group_sampler_keeps_group_ids_together(tmp_path: Path) -> None: + dataset = _make_toy_dataset(tmp_path) + sampler = GroupAwareBatchSampler(dataset, mode="full_group", batch_groups=1) + batch_indices = next(iter(sampler)) + records = [dataset[index] for index in batch_indices] + assert len({record.group_id for record in records}) == 1 + assert len(records) == 4 + + +def test_pair_sampler_returns_same_group_reward_ordered_pairs(tmp_path: Path) -> None: + dataset = _make_toy_dataset(tmp_path) + sampler = GroupAwareBatchSampler( + dataset, + mode="pairs", + batch_groups=3, + pair_count_per_group=2, + shuffle=False, + seed=3, + ) + batch_indices = next(iter(sampler)) + records = [dataset[index] for index in batch_indices] + assert batch_indices.pair_indices + for better_index, worse_index in batch_indices.pair_indices: + better = records[better_index] + worse = records[worse_index] + assert better.group_id == worse.group_id + assert better.reward.score > worse.reward.score + + +def test_collate_returns_tensors_and_metadata(tmp_path: Path) -> None: + dataset = _make_toy_dataset(tmp_path) + sampler = GroupAwareBatchSampler(dataset, mode="full_group", batch_groups=1) + records = [dataset[index] for index in next(iter(sampler))] + batch = collate_cil_records(records) + + assert batch["observations"].shape[0] == len(records) + assert batch["action_features"].shape[0] == len(records) + assert batch["effect_features"].shape[0] == len(records) + assert batch["rewards"].shape == (len(records),) + assert len(batch["instructions"]) == len(records) + assert len(batch["action_chunks"]) == len(records) + assert len(batch["effects"]) == len(records) + assert len(batch["group_ids"]) == len(records) + assert len(batch["candidate_types"]) == len(records) + assert len(batch["failures"]) == len(records) + assert "pair_indices" in batch + + +def test_zero_copy_collection_loads_disjoint_source_datasets(tmp_path: Path) -> None: + tasks = built_in_toy_tasks() + sources = [tmp_path / "source-a", tmp_path / "source-b"] + for source, task, seed in zip(sources, tasks[:2], (31, 47), strict=False): + generate_cil_dataset( + backend="toy", + tasks=[task], + out_dir=source, + num_states_per_task=2, + k=3, + seed=seed, + shard_size=8, + inline_observations=True, + ) + collection_dir = tmp_path / "collection" + write_cil_collection(collection_dir, sources, dataset_name="two-task-collection") + + dataset = CILDataset(collection_dir) + + assert len(dataset) == 12 + assert len(dataset.group_ids) == 4 + assert dataset.index.metadata["dataset_name"] == "two-task-collection" + assert dataset.index.metadata["task_count"] == 2 + assert {record.metadata["source_dataset"] for record in dataset} == { + str(source.resolve()) for source in sources + } diff --git a/tests/test_transfercritic.py b/tests/test_transfercritic.py new file mode 100644 index 0000000000000000000000000000000000000000..2775e86fb3e865086aa727e483c07d70363b8652 --- /dev/null +++ b/tests/test_transfercritic.py @@ -0,0 +1,120 @@ +from __future__ import annotations + +import pytest + +from dovla_cil.transfercritic.eval import compare_selection_strategies +from dovla_cil.transfercritic.labeling import make_utility_labels, toy_utility_value +from dovla_cil.transfercritic.schema import DataAtom, TransferContext +from dovla_cil.transfercritic.selection import greedy_marginal_selection + + +def _atom( + atom_id: str, + *, + score: float, + success: bool = False, + candidate_type: str = "near_miss", + task_id: str = "task_a", + cost: float = 1.0, +) -> DataAtom: + return DataAtom( + record_id=atom_id, + embedding=[score, 1.0 if success else 0.0, cost, 0.5], + candidate_type=candidate_type, + task_metadata={"task_id": task_id, "family": "pick"}, + reward_summary={ + "progress": score, + "score": score + (1.0 if success else 0.0), + "success": 1.0 if success else 0.0, + "regret": max(0.0, 1.0 - score), + }, + effect_summary={"moved_object_count": 1.0, "true_relation_count": score}, + cost=cost, + ) + + +def test_data_atom_and_context_schema_roundtrip() -> None: + atom = _atom("r1", score=0.5, success=True) + context = TransferContext( + benchmark_name="CausalStress", + task_family="pick", + target_objects=["red_mug"], + ood_factor="wrong_target", + validation_ref="val://small", + ) + + restored = DataAtom.from_dict(atom.to_dict()) + context_restored = TransferContext.from_dict(context.to_dict()) + + assert restored == atom + assert context_restored == context + assert restored.atom_id == "r1" + + +def test_greedy_selection_uses_score_over_cost() -> None: + atoms = [ + _atom("cheap_good", score=0.8, success=True, cost=1.0), + _atom("expensive_best", score=1.0, success=True, cost=3.0), + _atom("bad", score=0.1, success=False, cost=1.0), + ] + context = TransferContext(benchmark_name="CausalStress", task_family="pick") + + result = greedy_marginal_selection( + atoms, + context, + budget=2.0, + score_fn=lambda atom, _selected, _context: float(atom.reward_summary["score"]), + ) + + assert result.name == "transfercritic" + assert result.total_cost <= 2.0 + assert result.atom_ids[0] == "cheap_good" + assert "expensive_best" not in result.atom_ids + + +def test_toy_utility_labels_prefer_successful_useful_atoms() -> None: + context = TransferContext(benchmark_name="CausalStress", task_family="pick") + good = _atom("good", score=0.9, success=True, candidate_type="near_miss") + poor = _atom("poor", score=0.1, success=False, candidate_type="noop") + + labels = make_utility_labels([good, poor], context, method="toy_retraining_delta") + + assert labels[0].utility == toy_utility_value(good, context) + assert labels[0].utility > labels[1].utility + assert labels[0].metadata["approximate"] is True + + +def test_selection_experiments_return_all_baselines() -> None: + atoms = [ + _atom("a", score=0.6, task_id="task_a"), + _atom("b", score=0.9, success=True, task_id="task_b"), + _atom("c", score=0.2, task_id="task_a"), + ] + context = TransferContext(benchmark_name="CausalStress", task_family="pick") + + rows = compare_selection_strategies(atoms, context, budget=2.0, seed=0) + + assert {row["name"] for row in rows} == { + "random_subset", + "top_reward_subset", + "task_balanced_subset", + "transfercritic", + } + assert all(float(row["total_cost"]) <= 2.0 for row in rows) + + +def test_transfercritic_model_shapes() -> None: + torch = pytest.importorskip("torch") + from dovla_cil.transfercritic.model import SetConditionedTransferCritic, TransferCriticConfig + + config = TransferCriticConfig(atom_dim=4, set_dim=4, context_dim=4, hidden_dim=8) + model = SetConditionedTransferCritic(config) + atom = torch.randn(3, 4) + current_set = torch.zeros(3, 4) + context = torch.randn(3, 4) + + score = model(atom, current_set, context) + + assert score.shape == (3,) + score.mean().backward() + assert any(parameter.grad is not None for parameter in model.parameters()) diff --git a/tests/test_vlm_annotation.py b/tests/test_vlm_annotation.py new file mode 100644 index 0000000000000000000000000000000000000000..e5cde89dc9a18718535b2cdf3cd5a9359ee59a97 --- /dev/null +++ b/tests/test_vlm_annotation.py @@ -0,0 +1,186 @@ +from __future__ import annotations + +from dovla_cil.data.schema import ActionChunk, FailureInfo, RewardInfo, StructuredEffect +from dovla_cil.tasks.library import ToyTaskLibrary +from dovla_cil.vlm.annotation import VLMFailureAnnotator +from dovla_cil.vlm.client import VLMClient + + +def test_mock_vlm_annotation_works(monkeypatch, tmp_path) -> None: + monkeypatch.setenv("OPENCLAUDE_MOCK", "1") + task = ToyTaskLibrary().get_by_id("toy_pick_object_among_distractors") + local_failure = _failure("wrong_target") + annotator = VLMFailureAnnotator( + client=VLMClient(api_key="test-secret", model="mock-model"), + cache_path=tmp_path / "annotations.json", + ) + + annotated = annotator.annotate_failure( + task=task, + instruction=task.instruction, + action=_action(candidate_type="wrong_target", target="blue_mug"), + effect=_effect(), + reward=_reward(success=False), + local_failure=local_failure, + ) + + assert annotated.type == "wrong_target" + assert "Mock VLM explanation" in (annotated.language_explanation or "") + metadata = annotated.metadata["semantic_annotation"]["vlm_annotation"] + assert metadata["source"] == "vlm" + assert metadata["suggested_failure_type"] == "wrong_target" + assert "test-secret" not in str(annotated.to_dict()) + + +def test_annotation_cache_hit_avoids_client_call(tmp_path) -> None: + task = ToyTaskLibrary().get_by_id("toy_lift_can") + client = CountingAnnotationClient() + annotator = VLMFailureAnnotator(client=client, cache_path=tmp_path / "cache.json") + kwargs = { + "task": task, + "instruction": task.instruction, + "action": _action(candidate_type="no_motion", target="can"), + "effect": _effect(), + "reward": _reward(success=False), + "local_failure": _failure("no_motion"), + } + + first = annotator.annotate_failure(**kwargs) + second = annotator.annotate_failure(**kwargs) + + assert client.calls == 1 + assert first.language_explanation == second.language_explanation + assert second.metadata["semantic_annotation"]["vlm_annotation"]["cache_hit"] is True + + +def test_invalid_vlm_json_falls_back_to_local_explanation(tmp_path) -> None: + task = ToyTaskLibrary().get_by_id("toy_lift_can") + local_failure = _failure("no_motion") + annotator = VLMFailureAnnotator(client=InvalidAnnotationClient(), cache_path=tmp_path / "cache.json") + + annotated = annotator.annotate_failure( + task=task, + instruction=task.instruction, + action=_action(candidate_type="noop", target="can"), + effect=_effect(), + reward=_reward(success=False), + local_failure=local_failure, + ) + + assert annotated.type == local_failure.type + assert annotated.language_explanation == local_failure.language_explanation + metadata = annotated.metadata["semantic_annotation"]["vlm_annotation"] + assert metadata["source"] == "local_fallback" + assert metadata["fallback_error"] + + +def test_vlm_annotation_cannot_change_reward_or_local_failure_type(tmp_path) -> None: + task = ToyTaskLibrary().get_by_id("toy_pick_object_among_distractors") + reward = _reward(success=False) + local_failure = _failure("wrong_target") + annotator = VLMFailureAnnotator(client=OverridingAnnotationClient(), cache_path=tmp_path / "cache.json") + + annotated = annotator.annotate_failure( + task=task, + instruction=task.instruction, + action=_action(candidate_type="wrong_target", target="blue_mug"), + effect=_effect(), + reward=reward, + local_failure=local_failure, + ) + + assert reward.progress == 0.25 + assert reward.success is False + assert reward.terminal_success is False + assert annotated.type == "wrong_target" + assert annotated.metadata["semantic_annotation"]["suggested_failure_type"] == "success" + + +class CountingAnnotationClient: + def __init__(self) -> None: + self.calls = 0 + + def chat_json(self, system: str, user: str, schema_hint=None) -> dict: + del system, user, schema_hint + self.calls += 1 + return { + "failure_type": "no_motion", + "explanation": "The action produced no task-relevant motion.", + "avoidance_hint": "Choose an action that contacts the can.", + "confidence": 0.9, + } + + +class InvalidAnnotationClient: + def chat_json(self, system: str, user: str, schema_hint=None) -> dict: + del system, user, schema_hint + return {"failure_type": "no_motion", "confidence": "high"} + + +class OverridingAnnotationClient: + def chat_json(self, system: str, user: str, schema_hint=None) -> dict: + del system, user, schema_hint + return { + "failure_type": "success", + "explanation": "The VLM incorrectly claims this succeeded.", + "avoidance_hint": "No hint.", + "confidence": 1.0, + } + + +def _action(*, candidate_type: str, target: str) -> ActionChunk: + return ActionChunk( + representation="semantic", + values=[{"command": "grasp", "object": target}], + skill_type="grasp", + metadata={ + "candidate_type": candidate_type, + "intended_target": target, + "intended_relation": "grasped", + "difficulty": 0.3, + }, + ) + + +def _effect() -> StructuredEffect: + before = { + "objects": { + "red_mug": {"position": [0.0, 0.0, 0.03], "grasped": False, "lifted": False}, + "blue_mug": {"position": [0.4, 0.0, 0.03], "grasped": False, "lifted": False}, + }, + "robot": {"eef_position": [0.0, 0.0, 0.2], "gripper": "open", "held_object": None}, + } + after = { + "objects": { + "red_mug": {"position": [0.0, 0.0, 0.03], "grasped": False, "lifted": False}, + "blue_mug": {"position": [0.6, 0.0, 0.03], "grasped": False, "lifted": False}, + }, + "robot": {"eef_position": [0.4, 0.0, 0.2], "gripper": "closed", "held_object": None}, + } + return StructuredEffect( + object_pose_delta={"blue_mug": [0.2, 0.0, 0.0]}, + contact_events=[{"type": "touch", "object": "blue_mug"}], + relation_before={"grasped(red_mug)": False}, + relation_after={"grasped(red_mug)": False}, + grasp_success=False, + moved_objects=["blue_mug"], + symbolic_before=before, + symbolic_after=after, + ) + + +def _reward(*, success: bool) -> RewardInfo: + return RewardInfo( + progress=1.0 if success else 0.25, + success=success, + terminal_success=success, + dense_components={"partial": 0.25}, + ) + + +def _failure(failure_type: str) -> FailureInfo: + return FailureInfo( + type=failure_type, + symbolic_reason=f"local {failure_type}", + language_explanation=f"Local explanation for {failure_type}.", + ) diff --git a/tests/test_vlm_client.py b/tests/test_vlm_client.py new file mode 100644 index 0000000000000000000000000000000000000000..9ee215e6feef5ec4fb0986c8948d23dcb585ab04 --- /dev/null +++ b/tests/test_vlm_client.py @@ -0,0 +1,75 @@ +from __future__ import annotations + +import logging + +import pytest + +from dovla_cil.vlm.client import VLMClient, VLMParseError + + +def test_mock_mode_returns_deterministic_json(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("OPENCLAUDE_MOCK", "1") + monkeypatch.setenv("OPENCLAUDE_MODEL", "mock-model") + + client = VLMClient(api_key="test-secret") + first = client.chat_json("system", "user", schema_hint={"answer": "string"}) + second = client.chat_json("system", "user", schema_hint={"answer": "string"}) + + assert first == second + assert first["mock"] is True + assert first["model"] == "mock-model" + assert first["schema_keys"] == ["answer"] + + +def test_chat_json_parses_fenced_json() -> None: + class StaticTextClient(VLMClient): + def chat_text(self, system: str, user: str) -> str: + del system, user + return '```json\n{"answer": 7, "ok": true}\n```' + + client = StaticTextClient(api_key="test-secret", model="model") + assert client.chat_json("system", "user") == {"answer": 7, "ok": True} + + +def test_chat_json_extracts_first_json_object() -> None: + class StaticTextClient(VLMClient): + def chat_text(self, system: str, user: str) -> str: + del system, user + return 'Here is the result: {"answer": {"nested": true}}. Thanks.' + + client = StaticTextClient(api_key="test-secret", model="model") + assert client.chat_json("system", "user") == {"answer": {"nested": True}} + + +def test_parse_error_redacts_secret() -> None: + class BadTextClient(VLMClient): + def chat_text(self, system: str, user: str) -> str: + del system, user + return "not json test-secret" + + client = BadTextClient(api_key="test-secret", model="model") + with pytest.raises(VLMParseError) as exc_info: + client.chat_json("system", "user") + assert "test-secret" not in str(exc_info.value) + assert "***REDACTED***" in str(exc_info.value) + + +def test_repr_and_retry_logs_redact_secret( + caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch +) -> None: + class FailingClient(VLMClient): + def _request_once(self, *, messages, response_format): + del messages, response_format + raise RuntimeError(f"boom {self.api_key}") + + monkeypatch.setattr("dovla_cil.vlm.client.time.sleep", lambda _seconds: None) + client = FailingClient(api_key="test-secret", model="model", max_retries=1) + + assert "test-secret" not in repr(client) + with caplog.at_level(logging.WARNING, logger="dovla_cil.vlm.client"): + with pytest.raises(Exception) as exc_info: + client.chat_text("system", "user") + + assert "test-secret" not in caplog.text + assert "test-secret" not in str(exc_info.value) + assert "***REDACTED***" in caplog.text diff --git a/tests/test_vlm_task_generator.py b/tests/test_vlm_task_generator.py new file mode 100644 index 0000000000000000000000000000000000000000..e33ba25f27953b52f5c5c7b972f66eacf6a29600 --- /dev/null +++ b/tests/test_vlm_task_generator.py @@ -0,0 +1,78 @@ +from __future__ import annotations + +from pathlib import Path + +from dovla_cil.tasks.library import built_in_toy_tasks +from dovla_cil.tasks.validators import validate_task +from dovla_cil.utils.io import iter_jsonl +from dovla_cil.vlm.client import VLMClient +from dovla_cil.vlm.task_generator import ( + TaskGenerator, + default_task_generation_request, +) + + +class FakeTaskClient(VLMClient): + def __init__(self, responses): + super().__init__(api_key="fake", model="fake") + self.responses = list(responses) + self.calls = 0 + + def chat_json(self, system, user, schema_hint=None): + del system, user, schema_hint + self.calls += 1 + if not self.responses: + raise AssertionError("Unexpected VLM call") + return self.responses.pop(0) + + +def test_mock_vlm_returns_valid_tasks(monkeypatch) -> None: + monkeypatch.setenv("OPENCLAUDE_MOCK", "1") + request = default_task_generation_request(num_tasks=3) + tasks = TaskGenerator(VLMClient(api_key="fake", model="mock")).generate_tasks(request) + + assert len(tasks) == 3 + for task in tasks: + validate_task(task) + + +def test_invalid_generated_task_is_repaired() -> None: + valid = built_in_toy_tasks()[0].to_dict() + invalid = dict(valid) + invalid["success_predicates"] = [{"name": "grasped", "args": ["missing"]}] + client = FakeTaskClient([{"tasks": [invalid]}, valid]) + request = default_task_generation_request(num_tasks=1) + + tasks = TaskGenerator(client).generate_tasks(request) + + assert len(tasks) == 1 + assert tasks[0].task_id == valid["task_id"] + assert client.calls == 2 + + +def test_invalid_generated_task_is_skipped_when_repair_fails() -> None: + valid = built_in_toy_tasks()[0].to_dict() + invalid = dict(valid) + invalid["success_predicates"] = [{"name": "grasped", "args": ["missing"]}] + still_invalid = dict(invalid) + client = FakeTaskClient([{"tasks": [invalid]}, still_invalid]) + + tasks = TaskGenerator(client).generate_tasks(default_task_generation_request(num_tasks=1)) + + assert tasks == [] + assert client.calls == 2 + + +def test_script_runs_in_mock_mode_without_network(tmp_path: Path, monkeypatch) -> None: + from scripts.generate_tasks import main + + out = tmp_path / "tasks.jsonl" + monkeypatch.delenv("OPENCLAUDE_API_KEY", raising=False) + monkeypatch.delenv("OPENCLAUDE_MODEL", raising=False) + + result = main(["--num-tasks", "2", "--out", str(out), "--mock", "--seed", "5"]) + + assert result == 0 + rows = list(iter_jsonl(out)) + assert len(rows) == 2 + assert rows[0]["task_id"]