Port the ReBRAC implementation from JAX to PyTorch !
Only one file:
usage: rebrac_pytorch.py [-h] [OPTIONS]
╭─ options ────────────────────────────────────────────────────────────────────────────────────────╮
│ -h, --help show this help message and exit │
│ --project STR wandb params (default: ReBRAC) │
│ --group STR wandb params (default: rebrac-halfcheetah-medium-v2) │
│ --name STR wandb params (default: rebrac-new) │
│ --actor-learning-rate FLOAT │
│ model params (default: 0.001) │
│ --critic-learning-rate FLOAT │
│ model params (default: 0.001) │
│ --hidden-dim INT model params (default: 256) │
│ --actor-n-hiddens INT model params (default: 3) │
│ --critic-n-hiddens INT model params (default: 3) │
│ --gamma FLOAT model params (default: 0.99) │
│ --tau FLOAT model params (default: 0.005) │
│ --actor-bc-coef FLOAT model params (default: 0.001) │
│ --critic-bc-coef FLOAT model params (default: 0.01) │
│ --actor-ln, --no-actor-ln │
│ model params (default: False) │
│ --critic-ln, --no-critic-ln │
│ model params (default: True) │
│ --policy-noise FLOAT model params (default: 0.2) │
│ --noise-clip FLOAT model params (default: 0.5) │
│ --policy-freq INT model params (default: 2) │
│ --normalize-q, --no-normalize-q │
│ model params (default: True) │
│ --env-name STR training params (default: HalfCheetah-v5) │
│ --dataset-path STR training params (default: '~/.d4rl/datasets/halfcheetah_medium-v2.hdf5') │
│ --batch-size INT training params (default: 1024) │
│ --num-epochs INT training params (default: 1000) │
│ --num-updates-on-epoch INT │
│ training params (default: 1000) │
│ --normalize-reward, --no-normalize-reward │
│ training params (default: False) │
│ --normalize-states, --no-normalize-states │
│ training params (default: False) │
│ --eval-episodes INT evaluation params (default: 10) │
│ --eval-every INT evaluation params (default: 5) │
│ --train-seed INT general params (default: 0) │
│ --eval-seed INT general params (default: 42) │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯