reppo/slurm/run_reppo_dmc_prod.sh
ys1087@partner.kit.edu a02e258f1c seperate dmc setup...
2025-07-29 14:58:43 +02:00

56 lines
1.5 KiB
Bash
Executable File

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