Minimal Decision Transformer Implementation written in Jax (Flax). - View it on GitHub