run_nonlinear_experiments.sh

#!/bin/bash

# Activate conda environment
source ~/anaconda3/etc/profile.d/conda.sh
conda activate diffop

# Parse command line arguments
MAX_PARALLEL=${1:-4}  # Default to 4 parallel jobs, can be overridden: ./run_nonlinear_experiments.sh 8

# Environments: Cartpole-v0, Quadrotor-v0, Robotarm-v0

echo "Starting nonlinear experiments..."
echo "Environments: Cartpole-v0, Quadrotor-v0, Robotarm-v0"
echo "Seeds: 0, 1, 2, 3, 4"
echo "Max parallel jobs: $MAX_PARALLEL"
echo "=========================================="

# Function to wait for jobs when we reach max parallel limit
wait_for_jobs() {
    while [ $(jobs -r | wc -l) -ge $MAX_PARALLEL ]; do
        sleep 1
    done
}


run_ppo_experiment() {
    local env=$1
    local seed=$2
    echo "Running PPO on $env with seed $seed..."
    python diffop/experiments/train_PPO.py --env $env --seed $seed --project_name "DiffOP-Nonlinear"
}


run_diffop_experiment() {
    local env=$1
    local seed=$2
    local horizon=$3
    local lr=$4
    echo "Running DiffOP on $env with seed $seed, horizon $horizon, lr $lr..."
    python diffop/experiments/train_diffop.py --env $env --seed $seed --lr $lr --std 1e-2 --horizon $horizon --apply_horizon 1 --final_cost 0 --wandb_log --project_name "DiffOP-Nonlinear"
}

run_mpc_experiment() {
    local algorithm=$1
    local env=$2
    local seed=$3
    local horizon=$4
    echo "Running $algorithm on $env with seed $seed, horizon $horizon..."
    python diffop/experiments/train_$algorithm.py --env $env --seed $seed --horizon $horizon --wandb_log --project_name "DiffOP-Nonlinear"
}

# Run PPO experiments
# echo "Starting PPO experiments..."
# for env in "Cartpole-v0" "Quadrotor-v0" "Robotarm-v0"; do
#     for seed in {0..4}; do
#         run_ppo_experiment $env $seed
#         echo "Completed PPO $env seed $seed"
#         echo "----------------------------------------"
#     done
# done

echo "=========================================="
echo "Starting DiffOP Step experiments"

# Environment configurations
declare -A env_configs
env_configs["Cartpole-v0"]="6 0.2"
env_configs["Robotarm-v0"]="12 0.1"
env_configs["Quadrotor-v0"]="12 0.1"

# Run DiffOP experiments with configurations (parallel)
for env in "${!env_configs[@]}"; do
    IFS=' ' read -r horizon lr <<< "${env_configs[$env]}"
    echo "Starting DiffOP experiments for $env (horizon=$horizon, lr=$lr)..."
    
    for seed in {0..4}; do
        wait_for_jobs
        (
            run_diffop_experiment "$env" $seed $horizon $lr
            echo "Completed DiffOP $env seed $seed (horizon=$horizon, lr=$lr)"
        ) &
    done
done


echo "=========================================="
echo "Starting DiffOP Traj experiments"

# # DiffOP Traj experiments with horizon=0, lr=0.1 for all environments (parallel)
for env in "Cartpole-v0" "Quadrotor-v0" "Robotarm-v0"; do
    echo "Starting DiffOP Traj experiments for $env (horizon=0, lr=0.1)..."
    
    for seed in {0..4}; do
        wait_for_jobs
        (
            run_diffop_experiment "$env" $seed 0 0.1
            echo "Completed DiffOP Traj $env seed $seed (horizon=0, lr=0.1)"
        ) &
    done
done

# echo "=========================================="
# echo "Starting MPC experiments (MPCDPG & MPCTD)"

# Run MPCTD experiments with learning rate grid search
echo "Starting MPCTD experiments..."
for env in "${!env_configs[@]}"; do
    IFS=' ' read -r horizon lr <<< "${env_configs[$env]}"
    echo "Starting MPCTD for $env (horizon=$horizon)..."
    
    for seed in {0..4}; do
        for lr in 1e-3 1e-4 1e-5; do
            wait_for_jobs
            (
                python diffop/experiments/train_mpctd.py --env $env --seed $seed --horizon $horizon --lr $lr --wandb_log --project_name "DiffOP-Nonlinear"
                echo "Completed MPCTD $env seed $seed (horizon=$horizon, lr=$lr)"
            ) &
        done
    done
done

# Run MPCDPG experiments with learning rate and explore rate grid search
echo "Starting MPCDPG experiments..."
for env in "${!env_configs[@]}"; do
    IFS=' ' read -r horizon lr <<< "${env_configs[$env]}"
    echo "Starting MPCDPG for $env (horizon=$horizon)..."
    
    for seed in {0..4}; do
        for lr in 1e-2 1e-3 1e-4; do
            for explore_rate in 1e-1 1e-2 1e-3; do
                wait_for_jobs
                (
                    python diffop/experiments/train_mpcdpg.py --env $env --seed $seed --horizon $horizon --lr $lr --explore_rate $explore_rate --wandb_log --project_name "DiffOP-Nonlinear"
                    echo "Completed MPCDPG $env seed $seed (horizon=$horizon, lr=$lr, explore=$explore_rate)"
                ) &
            done
        done
    done
done

# Wait for all background jobs to complete
echo "Waiting for all experiments to complete..."
wait

echo "All experiments completed!"