Port the ReBRAC implementation from JAX to PyTorch !

Uasge

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)                                             │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯