#!/bin/bash
set -e

PY="/home/hongyi/miniconda3/envs/pwm/bin/python"
ROOT="/home/hongyi/PWM/EgoVLA_Release"
SCRIPT="$ROOT/dp_results/rl/chain_wm_rollout.py"
LORA="/home/hongyi/PWM/external/DiffSynth-Studio/models/train/Wan2.2-TI2V-5B-Robot_mixed_new/step-25000.safetensors"
OUT_BASE="$ROOT/dp_results/rl_runs"

export DIFFSYNTH_SKIP_DOWNLOAD=true
export DIFFSYNTH_MODEL_BASE_PATH="/home/hongyi/PWM/external/DiffSynth-Studio/models"

# Flip-Mug: max_horizon=400, 400/16=25 chunks
echo "========== Flip-Mug =========="
$PY $SCRIPT \
  --task Humanoid-Flip-Mug-v0 \
  --ckpt $ROOT/dp_results/runs/flip_mug_bc_10demos/checkpoints/latest.ckpt \
  --lora $LORA \
  --num_init_states 1 --n_chunks_per_init 1 --max_steps 400 \
  --wm_prompt "A humanoid robot flips a mug on a table with its dexterous hands." \
  --output_dir $OUT_BASE/flip_mug_10demos_mixed25k_debug/rollouts

# Pour-Balls: max_horizon=400, 400/16=25 chunks
echo "========== Pour-Balls =========="
$PY $SCRIPT \
  --task Humanoid-Pour-Balls-v0 \
  --ckpt $ROOT/dp_results/runs/pour_balls_bc_10demos/checkpoints/latest.ckpt \
  --lora $LORA \
  --num_init_states 1 --n_chunks_per_init 1 --max_steps 400 \
  --wm_prompt "A humanoid robot pours balls in cup into bowl with its dexterous hands." \
  --output_dir $OUT_BASE/pour_balls_10demos_mixed25k_debug/rollouts

# Stack-Can: max_horizon=300, 300/16=18 chunks
echo "========== Stack-Can =========="
$PY $SCRIPT \
  --task Humanoid-Stack-Can-v0 \
  --ckpt $ROOT/dp_results/runs/stack_can_bc_10demos/checkpoints/latest.ckpt \
  --lora $LORA \
  --num_init_states 1 --n_chunks_per_init 1 --max_steps 300 \
  --wm_prompt "A humanoid robot stack a can on the saucer with its dexterous hands." \
  --output_dir $OUT_BASE/stack_can_10demos_mixed25k_debug/rollouts

echo "========== All 3 debug rollouts done =========="