#!/bin/bash

# RATE
python3 src/train.py --wandb.project-name=RATE-Memory-Maze --wandb.wwandb=True --data.gamma=1.0 --data.path-to-dataset=data/Memory_Maze/ --training.learning-rate=0.0003 --training.lr-end-factor=0.1 --training.beta-1=0.9 --training.beta-2=0.95 --training.weight-decay=0.1 --training.batch-size=64 --training.warmup-steps=1_000 --training.final-tokens=500_000 --training.grad-norm-clip=1.0 --training.epochs=80 --training.ckpt-epoch=1 --training.online-inference=True --training.log-last-segment-loss-only=False --training.use-cosine-decay=True --training.context-length=30 --training.sections=3 --model.env-name=memory_maze --model.state-dim=3 --model.act-dim=6 --model.n-layer=6 --model.n-head=8 --model.n-head-ca=4 --model.d-model=64 --model.d-head=64 --model.d-inner=64 --model.dropout=0.5 --model.dropatt=0.2 --model.mem-len=360 --model.ext-len=0 --model.num-mem-tokens=15 --model.mem-at-end=True --model.mrv-act=relu --model.skip-dec-ffn=True --model.padding-idx=None --tensorboard-dir=runs/Memory_Maze/RATE --model-mode=RATE --arch-mode=TrXL --start-seed=1 --end-seed=3 --text=memory_maze --online-inference.use-argmax=False --online-inference.episode-timeout=1000 --online-inference.desired-return-1=18.1

# DT
python3 src/train.py --wandb.project-name=RATE-Memory-Maze --wandb.wwandb=True --data.gamma=1.0 --data.path-to-dataset=data/Memory_Maze/ --training.learning-rate=0.0003 --training.lr-end-factor=0.1 --training.beta-1=0.9 --training.beta-2=0.95 --training.weight-decay=0.1 --training.batch-size=64 --training.warmup-steps=1_000 --training.final-tokens=500_000 --training.grad-norm-clip=1.0 --training.epochs=80 --training.ckpt-epoch=1 --training.online-inference=True --training.log-last-segment-loss-only=False --training.use-cosine-decay=True --training.context-length=30 --training.sections=3 --model.env-name=memory_maze --model.state-dim=3 --model.act-dim=6 --model.n-layer=6 --model.n-head=8 --model.n-head-ca=0 --model.d-model=64 --model.d-head=64 --model.d-inner=64 --model.dropout=0.5 --model.dropatt=0.2 --model.mem-len=0 --model.ext-len=0 --model.num-mem-tokens=0 --model.mem-at-end=True --model.mrv-act=no_act --model.skip-dec-ffn=False --model.padding-idx=None --tensorboard-dir=runs/Memory_Maze/DT --model-mode=DT --arch-mode=TrXL --start-seed=1 --end-seed=3 --text=memory_maze --online-inference.use-argmax=False --online-inference.episode-timeout=1000 --online-inference.desired-return-1=18.1

# RMT
python3 src/train.py --wandb.project-name=RATE-Memory-Maze --wandb.wwandb=True --data.gamma=1.0 --data.path-to-dataset=data/Memory_Maze/ --training.learning-rate=0.0003 --training.lr-end-factor=0.1 --training.beta-1=0.9 --training.beta-2=0.95 --training.weight-decay=0.1 --training.batch-size=64 --training.warmup-steps=1_000 --training.final-tokens=500_000 --training.grad-norm-clip=1.0 --training.epochs=80 --training.ckpt-epoch=1 --training.online-inference=True --training.log-last-segment-loss-only=False --training.use-cosine-decay=True --training.context-length=30 --training.sections=3 --model.env-name=memory_maze --model.state-dim=3 --model.act-dim=6 --model.n-layer=6 --model.n-head=8 --model.n-head-ca=0 --model.d-model=64 --model.d-head=64 --model.d-inner=64 --model.dropout=0.5 --model.dropatt=0.2 --model.mem-len=0 --model.ext-len=0 --model.num-mem-tokens=15 --model.mem-at-end=True --model.mrv-act=no_act --model.skip-dec-ffn=False --model.padding-idx=None --tensorboard-dir=runs/Memory_Maze/RMT --model-mode=RMT --arch-mode=TrXL --start-seed=1 --end-seed=3 --text=memory_maze --online-inference.use-argmax=False --online-inference.episode-timeout=1000 --online-inference.desired-return-1=18.1

# TrXL
python3 src/train.py --wandb.project-name=RATE-Memory-Maze --wandb.wwandb=True --data.gamma=1.0 --data.path-to-dataset=data/Memory_Maze/ --training.learning-rate=0.0003 --training.lr-end-factor=0.1 --training.beta-1=0.9 --training.beta-2=0.95 --training.weight-decay=0.1 --training.batch-size=64 --training.warmup-steps=1_000 --training.final-tokens=500_000 --training.grad-norm-clip=1.0 --training.epochs=80 --training.ckpt-epoch=1 --training.online-inference=True --training.log-last-segment-loss-only=False --training.use-cosine-decay=True --training.context-length=30 --training.sections=3 --model.env-name=memory_maze --model.state-dim=3 --model.act-dim=6 --model.n-layer=6 --model.n-head=8 --model.n-head-ca=0 --model.d-model=64 --model.d-head=64 --model.d-inner=64 --model.dropout=0.5 --model.dropatt=0.2 --model.mem-len=300 --model.ext-len=0 --model.num-mem-tokens=0 --model.mem-at-end=True --model.mrv-act=no_act --model.skip-dec-ffn=False --model.padding-idx=None --tensorboard-dir=runs/Memory_Maze/TrXL --model-mode=TrXL --arch-mode=TrXL --start-seed=1 --end-seed=3 --text=memory_maze --online-inference.use-argmax=False --online-inference.episode-timeout=1000 --online-inference.desired-return-1=18.1