#!/bin/bash

if [ -z "$1" ]; then
    echo "Enter algorithm type"
    exit 0
else
    ALG=$1
    shift
fi

if [ $ALG == "a2c" ]; then
    ALG="python -m baselines.a2c.run_atari"
elif [ $ALG == "ppo" ]; then
    ALG="srun python -m baselines.ppo1.run_atari"
elif [ $ALG == "acktr" ]; then
    ALG="python -m baselines.acktr.run_atari"
else
    echo "Unsupported algorithm: "$ALG
    exit 0
fi

if [ -z "$1" ]; then
    echo "Enter environment name"
    exit 0
else
    ENV=$1
    shift
fi

if [ -z "$1" ]; then
    echo "Enter number of training steps (in millions)"
    exit 0
else
    FRAMES=$1
    shift
fi

if [ -z "$1" ]; then
    echo "Enter output directory"
    exit 0
else
    OUTPUT=$1
    shift
fi

$ALG --env $ENV --million_frames $FRAMES --outputDir $OUTPUT "$@"
