Add batch-size option

sony · Aug 28, 2023 · 0028b64 · 0028b64
1 parent 4adaa11
commit 0028b64
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 5 deletions.
diff --git a/bin/evaluate_algorithm b/bin/evaluate_algorithm
@@ -132,6 +132,8 @@ do
             SAVE_DIR="$1";;
         "--num_seeds" )
             NUM_SEEDS="$1";;
+        "--batch_size" )
+            BATCH_SIZE="$1";;
         "-h" )
             echo ""
             echo "[Usage] : ./evaluate_algo
@@ -142,6 +144,7 @@ do
                 --env_end_index <env_end_index>
                 --envs <env1> <env2> <env3>
                 --save_dir <save_dir>
+                --batch_size <batch_size>
 
                 e.g. ./evaluate_algo --algo_name dqn --base_env_name atari --env_start_index 0 --env_end_index 1 --save_dir sample
                 or if you want directly specify the training envs
@@ -183,5 +186,9 @@ do
         ENV_NAME=${DELAYED_MUJOCO_ENV_LIST[$INDEX]}
     fi
     echo "Start running training for: " ${ENV_NAME}
-    ${ROOT_DIR}/bin/train_with_seeds "${REPRODUCTION_CODE_DIR}/${ALGO_NAME}_reproduction.py" $GPU_ID $ENV_NAME $SAVE_DIR $NUM_SEEDS &
+    if [ -n "$BATCH_SIZE" ]; then
+        ${ROOT_DIR}/bin/train_with_seeds "${REPRODUCTION_CODE_DIR}/${ALGO_NAME}_reproduction.py" $GPU_ID $ENV_NAME $SAVE_DIR $NUM_SEEDS $BATCH_SIZE &
+    else
+        ${ROOT_DIR}/bin/train_with_seeds "${REPRODUCTION_CODE_DIR}/${ALGO_NAME}_reproduction.py" $GPU_ID $ENV_NAME $SAVE_DIR $NUM_SEEDS &
+    fi
 done
diff --git a/bin/train_with_seeds b/bin/train_with_seeds
@@ -13,8 +13,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-if [ $# -ne 5 ]; then
-    echo "usage: $0 <script_file_name> <gpu_id> <env> <save_dir> <num_seeds>"
+if [ $# -lt 5 ] && [ $# -gt 6 ]; then
+    echo "usage: $0 <script_file_name> <gpu_id> <env> <save_dir> <num_seeds> <batch_size>"
     exit 1
 fi
 RESULTDIR="$4/$3_results"
@@ -23,11 +23,19 @@ NUM_SEEDS=$5
 if [ $NUM_SEEDS -eq 3 ]; then
     for seed in 1 10 100
     do
-        python $1 --gpu $2 --seed $seed --env $3 --save-dir $4
+        if [ $# -eq 5 ]; then
+            python $1 --gpu $2 --seed $seed --env $3 --save-dir $4
+        else
+            python $1 --gpu $2 --seed $seed --env $3 --save-dir $4 --batch-size $6
+        fi
     done
 else
     for seed in $(seq 1 $NUM_SEEDS);
     do
-        python $1 --gpu $2 --seed $seed --env $3 --save-dir $4
+        if [ $# -eq 5 ]; then
+            python $1 --gpu $2 --seed $seed --env $3 --save-dir $4
+        else
+            python $1 --gpu $2 --seed $seed --env $3 --save-dir $4 --batch-size $6
+        fi
     done
 fi