Stars
Simple single-file baselines for Q-Learning in pure-GPU setting
Reinforcement learning on general 2D physics environments in JAX
Lightweight Cluster/Cloud VM Job Management 🚀
🕹️ A diverse suite of scalable reinforcement learning environments in JAX
[ICLR 2024] Closing the Gap between TD Learning and Supervised Learning - A Generalisation Point of View.
🏛️A research-friendly codebase for fast experimentation of single-agent reinforcement learning in JAX • End-to-End JAX RL
Maximal Update Parametrization (μP) with Flax & Optax.
Github action for syncing other repositories (templates) with current repository. Any git provider like GitHub (enterprise), GitLab, Gittea,.. are supported for the source repository
TORAX: Tokamak transport simulation in JAX
Octo is a transformer-based robot policy trained on a diverse mix of 800k robot trajectories.
😎 A curated list of awesome GitHub Profile which updates in real time
Flexible and scalable template based on PyTorch Lightning + Hydra. Efficient workflow and reproducibility for rapid ML experiments.
Mastering Diverse Domains through World Models
A logical, reasonably standardized, but flexible project structure for doing and sharing data science work.
A curated list of resources about generative flow networks (GFlowNets).
Transformer with Mu-Parameterization, implemented in Jax/Flax. Supports FSDP on TPU pods.
(Crafter + NetHack) in JAX. ICML 2024 Spotlight.
Official implementation of the δ-model presented in the ICML 2024 paper "A Distributional Analogue to the Successor Representation".
GFlowNet library specialized for graph & molecular data
A collection of utilities for machine learning experiments.
Selects tests affected by changed files. Executes the right tests first. Continuous test runner when used with pytest-watch.
Fast and simple implementation of RL algorithms, designed to run fully on GPU.