#!/bin/bash
#SBATCH --job-name=reppo_dmc_prod
#SBATCH --account=hk-project-p0022232
#SBATCH --partition=accelerated
#SBATCH --gres=gpu:1
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --cpus-per-task=8
#SBATCH --time=24:00:00
#SBATCH --mem=32G
#SBATCH --output=logs/reppo_dmc_prod_%j.out
#SBATCH --error=logs/reppo_dmc_prod_%j.err

# Load required modules
module load devel/cuda/12.4

# Set environment variables
export WANDB_MODE=online
export WANDB_PROJECT=reppo_dmc_production
export WANDB_API_KEY=01fbfaf5e2f64bedd68febedfcaa7e3bbd54952c
export WANDB_ENTITY=dominik_roth

# Change to project directory
cd /hkfs/home/project/hk-project-robolear/ys1087/Projects/reppo

# Activate virtual environment
source .venv/bin/activate

# Run DMC experiment
echo "Starting REPPO production run with DMC..."
echo "Job ID: $SLURM_JOB_ID"
echo "Node: $SLURM_NODELIST"
echo "GPU: $CUDA_VISIBLE_DEVICES"

# Environment name passed as variable
ENV_NAME=${ENV_NAME:-CartpoleBalance}
SEED=${SEED:-0}

echo "Environment: $ENV_NAME"
echo "Seed: $SEED"

# Run the experiment with full 50M steps
python reppo_alg/jaxrl/reppo.py \
    env=mjx_dmc \
    env.name=$ENV_NAME \
    hyperparameters.num_envs=1024 \
    hyperparameters.num_steps=128 \
    hyperparameters.num_mini_batches=128 \
    hyperparameters.num_epochs=4 \
    hyperparameters.total_time_steps=50000000 \
    seed=$SEED \
    wandb.mode=online \
    wandb.entity=$WANDB_ENTITY \
    wandb.project=$WANDB_PROJECT

echo "Training completed!"