JaxUED: A simple and useable UED library in Jax

• URL: http://arxiv.org/abs/2403.13091v1
• Date: Tue, 19 Mar 2024 18:40:50 GMT
• Title: JaxUED: A simple and useable UED library in Jax
• Authors: Samuel Coward, Michael Beukman, Jakob Foerster,
• Abstract summary: We present JaxUED, an open-source library providing minimal dependency implementations of modern Unsupervised Environment Design (UED) algorithms in Jax. Inspired by CleanRL, we provide fast, clear, understandable, and easily modifiable implementations, with the aim of accelerating research into UED.
• Score: 1.5821811088000381
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: We present JaxUED, an open-source library providing minimal dependency implementations of modern Unsupervised Environment Design (UED) algorithms in Jax. JaxUED leverages hardware acceleration to obtain on the order of 100x speedups compared to prior, CPU-based implementations. Inspired by CleanRL, we provide fast, clear, understandable, and easily modifiable implementations, with the aim of accelerating research into UED. This paper describes our library and contains baseline results. Code can be found at https://github.com/DramaCow/jaxued.

Related papers

• JaxDecompiler: Redefining Gradient-Informed Software Design [0.0]
  JaxDecompiler is a tool that transforms any JAX function into an editable Python code. This article introduces JaxDecompiler, a tool that transforms any JAX function into an editable Python code.
  arXiv  Detail & Related papers  (2024-03-14T20:32:31Z)
• JaxMARL: Multi-Agent RL Environments and Algorithms in JAX [105.343918678781]
  We present JaxMARL, the first open-source, Python-based library that combines GPU-enabled efficiency with support for a large number of commonly used MARL environments. Our experiments show that, in terms of wall clock time, our JAX-based training pipeline is around 14 times faster than existing approaches. We also introduce and benchmark SMAX, a JAX-based approximate reimplementation of the popular StarCraft Multi-Agent Challenge.
  arXiv  Detail & Related papers  (2023-11-16T18:58:43Z)
• QDax: A Library for Quality-Diversity and Population-based Algorithms with Hardware Acceleration [3.8494302715990845]
  QDax is an open-source library with a streamlined and modular API for Quality-Diversity (QD) optimization algorithms in Jax. The library serves as a versatile tool for optimization purposes, ranging from black-box optimization to continuous control.
  arXiv  Detail & Related papers  (2023-08-07T15:29:44Z)
• LeanDojo: Theorem Proving with Retrieval-Augmented Language Models [72.54339382005732]
  Large language models (LLMs) have shown promise in proving formal theorems using proof assistants such as Lean. Existing methods are difficult to reproduce or build on, due to private code, data, and compute requirements. This paper introduces LeanDojo: an open-source Lean toolkit consisting of toolkits, data, models. We develop ReProver: an LLM-based prover augmented with retrieval for selecting premises from a vast math library.
  arXiv  Detail & Related papers  (2023-06-27T17:05:32Z)
• JaxPruner: A concise library for sparsity research [46.153423603424]
  JaxPruner is an open-source library for sparse neural network research. It implements popular pruning and sparse training algorithms with minimal memory and latency overhead.
  arXiv  Detail & Related papers  (2023-04-27T10:45:30Z)
• SequeL: A Continual Learning Library in PyTorch and JAX [50.33956216274694]
  SequeL is a library for Continual Learning that supports both PyTorch and JAX frameworks. It provides a unified interface for a wide range of Continual Learning algorithms, including regularization-based approaches, replay-based approaches, and hybrid approaches. We release SequeL as an open-source library, enabling researchers and developers to easily experiment and extend the library for their own purposes.
  arXiv  Detail & Related papers  (2023-04-21T10:00:22Z)
• Bringing UMAP Closer to the Speed of Light with GPU Acceleration [28.64858826371568]
  We show a number of techniques that can be used to make a faster and more faithful GPU version of UMAP. Many of these design choices/lessons are general purpose and may inform the conversion of other graph and manifold learning algorithms to use GPU.
  arXiv  Detail & Related papers  (2020-08-01T19:35:56Z)
• PolyDL: Polyhedral Optimizations for Creation of High Performance DL primitives [55.79741270235602]
  We present compiler algorithms to automatically generate high performance implementations of Deep Learning primitives. We develop novel data reuse analysis algorithms using the polyhedral model. We also show that such a hybrid compiler plus a minimal library-use approach results in state-of-the-art performance.
  arXiv  Detail & Related papers  (2020-06-02T06:44:09Z)
• Torch-Struct: Deep Structured Prediction Library [138.5262350501951]
  We introduce Torch-Struct, a library for structured prediction. Torch-Struct includes a broad collection of probabilistic structures accessed through a simple and flexible distribution-based API.
  arXiv  Detail & Related papers  (2020-02-03T16:43:02Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.