NAVIX: Scaling MiniGrid Environments with JAX

• URL: http://arxiv.org/abs/2407.19396v1
• Date: Sun, 28 Jul 2024 04:39:18 GMT
• Title: NAVIX: Scaling MiniGrid Environments with JAX
• Authors: Eduardo Pignatelli, Jarek Liesen, Robert Tjarko Lange, Chris Lu, Pablo Samuel Castro, Laura Toni,
• Abstract summary: We introduce NAVIX, a re-implementation of MiniGrid in JAX. NAVIX achieves over 200 000x speed improvements in batch mode, supporting up to 2048 agents in parallel on a single Nvidia A100 80 GB. This reduces experiment times from one week to 15 minutes, promoting faster design and more scalable RL model development.
• Score: 17.944645332888335
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: As Deep Reinforcement Learning (Deep RL) research moves towards solving large-scale worlds, efficient environment simulations become crucial for rapid experimentation. However, most existing environments struggle to scale to high throughput, setting back meaningful progress. Interactions are typically computed on the CPU, limiting training speed and throughput, due to slower computation and communication overhead when distributing the task across multiple machines. Ultimately, Deep RL training is CPU-bound, and developing batched, fast, and scalable environments has become a frontier for progress. Among the most used Reinforcement Learning (RL) environments, MiniGrid is at the foundation of several studies on exploration, curriculum learning, representation learning, diversity, meta-learning, credit assignment, and language-conditioned RL, and still suffers from the limitations described above. In this work, we introduce NAVIX, a re-implementation of MiniGrid in JAX. NAVIX achieves over 200 000x speed improvements in batch mode, supporting up to 2048 agents in parallel on a single Nvidia A100 80 GB. This reduces experiment times from one week to 15 minutes, promoting faster design iterations and more scalable RL model development.

Related papers

• PCGRL+: Scaling, Control and Generalization in Reinforcement Learning Level Generators [2.334978724544296]
  Procedural Content Generation via Reinforcement Learning (PCGRL) has been introduced as a means by which controllable designer agents can be trained. PCGRL offers a unique set of affordances for game designers, but it is constrained by the compute-intensive process of training RL agents. We implement several PCGRL environments in Jax so that all aspects of learning and simulation happen in parallel on the GPU.
  arXiv  Detail & Related papers  (2024-08-22T16:30:24Z)
• SwiftRL: Towards Efficient Reinforcement Learning on Real Processing-In-Memory Systems [21.133750045141802]
  Reinforcement Learning (RL) trains agents to learn optimal behavior by maximizing reward signals from experience datasets. To overcome this, SwiftRL explores Processing-In-Memory (PIM) architectures to accelerate RL workloads. We achieve near-linear performance scaling by implementing RL algorithms like Tabular Q-learning and SARSA on UPMEM PIM systems.
  arXiv  Detail & Related papers  (2024-05-07T02:54: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)
• SRL: Scaling Distributed Reinforcement Learning to Over Ten Thousand Cores [13.948640763797776]
  We present a novel abstraction on the dataflows of RL training, which unifies diverse RL training applications into a general framework. We develop a scalable, efficient, and distributed RL system called ReaLly scalableRL, which allows efficient and massively parallelized training. SRL is the first in the academic community to perform RL experiments at a large scale with over 15k CPU cores.
  arXiv  Detail & Related papers  (2023-06-29T05:16:25Z)
• Rethinking Closed-loop Training for Autonomous Driving [82.61418945804544]
  We present the first empirical study which analyzes the effects of different training benchmark designs on the success of learning agents. We propose trajectory value learning (TRAVL), an RL-based driving agent that performs planning with multistep look-ahead. Our experiments show that TRAVL can learn much faster and produce safer maneuvers compared to all the baselines.
  arXiv  Detail & Related papers  (2023-06-27T17:58:39Z)
• EnvPool: A Highly Parallel Reinforcement Learning Environment Execution Engine [69.47822647770542]
  parallel environment execution is often the slowest part of the whole system but receives little attention. With a curated design for paralleling RL environments, we have improved the RL environment simulation speed across different hardware setups. On a high-end machine, EnvPool achieves 1 million frames per second for the environment execution on Atari environments and 3 million frames per second on MuJoCo environments.
  arXiv  Detail & Related papers  (2022-06-21T17:36:15Z)
• ElegantRL-Podracer: Scalable and Elastic Library for Cloud-Native Deep Reinforcement Learning [141.58588761593955]
  We present a library ElegantRL-podracer for cloud-native deep reinforcement learning. It efficiently supports millions of cores to carry out massively parallel training at multiple levels. At a low-level, each pod simulates agent-environment interactions in parallel by fully utilizing nearly 7,000 GPU cores in a single GPU.
  arXiv  Detail & Related papers  (2021-12-11T06:31:21Z)
• RAPID-RL: A Reconfigurable Architecture with Preemptive-Exits for Efficient Deep-Reinforcement Learning [7.990007201671364]
  We propose a reconfigurable architecture with preemptive exits for efficient deep RL (RAPID-RL) RAPID-RL enables conditional activation of preemptive layers based on the difficulty level of inputs. We show that RAPID-RL incurs 0.34x (0.25x) number of operations (OPS) while maintaining performance above 0.88x (0.91x) on Atari (Drone navigation) tasks.
  arXiv  Detail & Related papers  (2021-09-16T21:30:40Z)
• FNAS: Uncertainty-Aware Fast Neural Architecture Search [54.49650267859032]
  Reinforcement learning (RL)-based neural architecture search (NAS) generally guarantees better convergence yet suffers from the requirement of huge computational resources. We propose a general pipeline to accelerate the convergence of the rollout process as well as the RL process in NAS. Experiments on the Mobile Neural Architecture Search (MNAS) search space show the proposed Fast Neural Architecture Search (FNAS) accelerates standard RL-based NAS process by 10x.
  arXiv  Detail & Related papers  (2021-05-25T06:32:52Z)
• The NetHack Learning Environment [79.06395964379107]
  We present the NetHack Learning Environment (NLE), a procedurally generated rogue-like environment for Reinforcement Learning research. We argue that NetHack is sufficiently complex to drive long-term research on problems such as exploration, planning, skill acquisition, and language-conditioned RL. We demonstrate empirical success for early stages of the game using a distributed Deep RL baseline and Random Network Distillation exploration.
  arXiv  Detail & Related papers  (2020-06-24T14:12:56Z)

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.