#!/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"