Unofficial implemenation of Online Decision Transformer for MuJoCo robotics experiments. Please reach out if mistakes are found or you have any questions.
This codebase is based off the original codebase from Decision Transformer: Reinforcement Learning via Sequence Modeling.
As with the original codebase, experiments require MuJoCo. Follow the instructions in the mujoco-py repo to install. Then, dependencies can be installed with the following command:
conda env create -f environment.ymlDatasets are stored in the gym/data directory.
Install the D4RL repo, following the instructions there.
Then, run the following script in order to download the datasets and save them in our format:
python download_d4rl_datasets.pyAlternatively, the environment can be setup using Docker with the attached Dockerfile.
There are some small implementation detail differences and other assumed implementation detials that may be different.
However, we use $$\mathbf{tanh}(N(\mu\theta(s{-K,t}, g{-K,t}), \Sigma{\theta}(s{-K,t}, g{-K,t})))$$
We first sample from a normal distribution with parameters outputted by the transformer policy and then squash the sampled value to the action space using tanh.
There is no analytical form for the entropy of this distribution, but we can evaluate the (log) probability of a sample. We use a Monte Carlo estimate over a batch, similarly to the paper. Specifically we also sample $k$ actions for each transition within the batch to more closely approximate the entropy.
Emperically, we find that using tanh leads to better performance, but this change can be ignored by omitting --stochastic_tanh.
We train policies by first performing offline pretraining, and then perform online finetuning seperately.
The below command runs offline pretraining with hopper:
python experiment.py --env hopper --dataset medium --model_type dt --num_eval_episodes=50 --max_iters=5 --num_steps_per_iter=1000 --stochastic --use_action_means --learning_rate=1e-4 --embed_dim=512 --weight_decay=5e-4 --K=20 --remove_pos_embs --n_layer=4 --n_head=4 --batch_size=256 --eval_context=5 --stochastic_tanhpython experiment.py --env hopper --dataset medium --model_type dt --pretrained_model=./models/hopper/dt_gym-experiment-hopper-medium-506105.pt --stochastic --use_action_means --online_training --eval_context=5 --K=20 --batch_size=256 --num_steps_per_iter=300 --max_iters=200 --num_eval_episodes=50 --device=cuda:2 --target_entropy --stochastic_tanhResults can be logged by using --log_to_wandb=True.
Experimental wandb results can be found for online learning here for hopper. This implementation seems to converge after ~60-70 exploration rounds instead of ~250 as shown in the paper.