#!/bin/bash #SBATCH --job-name=reppo_brax #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=04:00:00 #SBATCH --mem=24G #SBATCH --output=logs/reppo_brax_%j.out #SBATCH --error=logs/reppo_brax_%j.err # Load required modules module load devel/cuda/12.4 # Set environment variables export WANDB_MODE=online export WANDB_PROJECT=reppo_brax # Change to project directory cd /hkfs/home/project/hk-project-robolear/ys1087/Projects/reppo # Activate virtual environment source .venv/bin/activate # Note: Ensure WANDB_API_KEY and WANDB_ENTITY are set before running # Run REPPO with Brax environment echo "Starting REPPO training with Brax..." echo "Job ID: $SLURM_JOB_ID" echo "Node: $SLURM_NODELIST" echo "GPU: $CUDA_VISIBLE_DEVICES" # Default environment: ant (can be overridden) ENV_NAME=${ENV_NAME:-ant} EXPERIMENT_TYPE=${EXPERIMENT_TYPE:-mjx_dmc_small_data} echo "Environment: $ENV_NAME" echo "Experiment type: $EXPERIMENT_TYPE" # Run the experiment python reppo_alg/jaxrl/reppo.py \ env=brax \ env_name=$ENV_NAME \ experiment_override=$EXPERIMENT_TYPE \ wandb.mode=online \ wandb.entity=${WANDB_ENTITY} \ wandb.project=$WANDB_PROJECT \ wandb.name="reppo_${ENV_NAME}_${EXPERIMENT_TYPE}_${SLURM_JOB_ID}" echo "Training completed!"