PaperlayerView exact revision on arXiv

Source-linked AI summary

MuJoCo-Drones-Gym: A GPU-Accelerated Multi-Drone Simulator for Control and Reinforcement Learning

Manan Tayal

arXiv:2606.08039v1cs.RO

TL;DR

Existing quadcopter learning environments limit the combination of physical fidelity, multi-agent support, and deep-RL throughput. MuJoCo-Drones-Gym addresses this with a modular Gymnasium- and PettingZoo-compatible MuJoCo simulator whose GPU-vectorized backend supports batched JAX-native training and broad task interfaces.

  • Problem

    Existing quadcopter learning environments face limitations in physical fidelity, multi-agent support, GPU vectorization, and compatibility with modern Gymnasium and PettingZoo APIs.

  • Method

    The paper introduces a modular MuJoCo-based multi-drone environment with configurable physics, actions, observations, wrappers, task templates, and an MJX GPU-vectorized backend.

  • Results

    MuJoCo-Drones-Gym provides seven task templates, six physics modes, three drone models, composable wrappers, and batched JAX-native training with a default 4096 parallel environments.

  • Takeaways & Limitations

    The environment preserves gym-pybullet-drones design choices so controllers and policies can port between the two with minimal effort.

  • Takeaways & Limitations

    The MJX backend currently supports only rigid-body MuJoCo physics and a limited task, observation, and action subset, omitting aerodynamic effects.

Abstract

from arXiv · show

Robotic simulators are a cornerstone of modern research in aerial robotics, serving both as a vehicle for the development of new control algorithms and as the data source for training reinforcement learning (RL) policies. Yet, existing quadcopter learning environments often face a trade-off between physical fidelity, multi-agent support, and the throughput required by modern deep RL pipelines. In this paper, we present MuJoCo-Drones-Gym, an open-source Gymnasium-compatible multi-drone environment built on top of the MuJoCo physics engine. MuJoCo-Drones-Gym supports an arbitrary number of Bitcraze Crazyflie 2.x nano-quadcopters and exposes a modular API for selecting (i)~the physics model (rigid-body MuJoCo, explicit Python dynamics, or any subset of ground effect, blade drag, and inter-drone downwash), (ii)~the action interface (per-motor RPMs, collective normalized thrust, velocity setpoints, or PID waypoint commands), and (iii)~the observation space (kinematic state vectors, RGB / depth / segmentation cameras, or neighbourhood adjacency information). A PettingZoo ParallelEnv wrapper enables drop-in multi-agent reinforcement learning, while a suite of seven task environments, hover, velocity tracking, multi-drone hover, waypoint navigation, formation flight, gate racing, and a generic multi-agent template, demonstrates the breadth of the interface. We describe the environment design, the underlying physics and quadcopter dynamics, and illustrate its use through control and learning examples that mirror those of the closely related gym-pybullet-drones project, while taking advantage of MuJoCo's improved contact handling, rendering, and parallelizability.

1 Introduction

MuJoCo-Drones-Gym is introduced as an open-source, Gymnasium-compatible multi-quadcopter environment that extends gym-pybullet-drones with MuJoCo physics, multi-agent support, and GPU-vectorized simulation. Its contributions include selectable physics and controllers, seven task environments, composable wrappers, multiple drone models, and integrations for modern RL APIs.

  • Core environment: MuJoCo-Drones-Gym ports the gym-pybullet-drones design philosophy to MuJoCo as an open-source multi-quadcopter learning environment.It is presented as the paper’s central package contribution.
  • Physics and control: Six selectable physics modes cover pure MuJoCo, ground effect, drag, downwash, all-combined, and explicit-integration fallback dynamics.Rotor forcing uses xfrc applied-based actuation.
  • Task suite: Seven task environments are provided: HoverAviary, VelocityAviary, MultiHoverAviary, FlyThroughAviary, FormationAviary, RaceAviary, and MultiAgentAviary.The suite spans hovering, velocity control, navigation, formation, racing, and generic multi-agent use cases.
  • Multi-agent RL: A PettingZoo ParallelEnv wrapper enables native multi-agent reinforcement learning on the same physics engine.The package also aligns with modern Gymnasium and PettingZoo APIs that have replaced OpenAI Gym as community standards.
  • Extensibility and assets: Additional contributions include cascaded PID controllers, wind and obstacle wrappers with automatic curriculum, three system-identified drone models, and vendored Crazyflie meshes.The models are CF2X, CF2P, and a 250 g RACE class; bundled meshes avoid requiring the full MuJoCo Menagerie.
  • GPU acceleration: GPU vectorization uses MJXVectorAviary to run thousands of parallel environments through one jax.vmap / jit call, with a Gymnasium VectorEnv wrapper for numpy-based RL libraries.The wrapper targets Stable-Baselines3 and other numpy-based reinforcement-learning libraries.

2 Background and Related Work

MuJoCo-drones-gym targets lightweight quadrotor simulation while addressing gym-pybullet-drones’ limitations through a modern MuJoCo/MJX backend, native GPU vectorization, and first-class multi-agent support. It occupies this niche alongside related environments with distinct trade-offs in physics fidelity, deployment, rendering, and hardware integration.

  • Related simulator landscape: MuJoCo-drones-gym targets the lightweight simulator niche while filling gym-pybullet-drones’ gap with a modern physics backend and native GPU vectorization.It remains lightweight while adopting MuJoCo/MJX.
  • Related simulator landscape: gym-pybullet-drones provides the direct API, aerodynamic-model, and PID ancestors, enabling controllers and policies to port with minimal effort.Its models include ground effect, blade drag, and downwash.
  • Related simulator landscape: safe-control-gym and CrazySwarm2 offer simpler physics and lack a first-class multi-agent or GPU-vectorized interface.safe-control-gym emphasizes classical and safe learning-based control, whereas CrazySwarm2 emphasizes ROS 2 firmware-in-the-loop and real-hardware integration.
  • Related simulator landscape: Flightmare combines Unity rendering with fast custom dynamics for single-agent Gym-style control but lacks a native multi-agent API and is heavier to install.It is not a pip-only package.
  • Related simulator landscape: OmniDrones targets large-scale GPU reinforcement learning with high throughput but depends tightly on Isaac Sim/Isaac Lab and proprietary GPU drivers.This restricts deployment to a narrower platform set than a pure JAX/MJX pipeline.
  • MuJoCo-drones-gym positioning: MuJoCo-drones-gym combines lightweight installation, pure Python, Gymnasium compatibility, MuJoCo/MJX, PettingZoo ParallelEnv, and richer task, controller, and disturbance support.A detailed feature comparison is deferred to Table 7 in Section 5.

3 Methodology

The methodology defines the simulator’s dynamics, physics and drone-model options, modular API, optional wrappers, and GPU-vectorized MJX backend.

  • 3 Methodology: The methodology covers quadrotor dynamics and aerodynamic models, six physics modes, three drone models, software architecture, and public action and observation spaces.It also introduces optional wind, obstacle, curriculum, and domain-randomization wrappers.
  • 3 Methodology: A GPU-vectorized MJX backend completes the simulator’s methodological design.

3.1 Quadrotor Dynamics

The simulator models each drone with a 13-dimensional state and four motor-speed controls, using quadratic thrust and reaction-torque relationships to generate force and moments. Three independently configurable effects—ground effect, blade drag, and inter-drone downwash—extend the dynamics model.

  • State and control: Each drone state is x = [p, q, v, ω]⊤ ∈ R13, comprising position, body-to-world quaternion, linear velocity, and angular velocity.The quaternion follows MuJoCo’s [w, x, y, z] convention.
  • State and control: The control input is the four-motor speed vector n = [n0, n1, n2, n3]⊤, expressed in RPM before appropriate conversion.Motor speeds are converted to radians per second for the dynamics.
  • Motor-generated forces and moments: Per-motor thrust and reaction torque are quadratic in motor speed, with alternating reaction-torque signs providing yaw control.For X-configurations, roll and pitch torques are computed from motor thrusts; the +-configuration uses a different motor-arm arrangement.
  • Optional aerodynamic effects: Three optional effects can be toggled independently: low-altitude ground effect, motor-speed-scaled blade drag, and downward inter-drone downwash.Ground effect is clipped near zero altitude and gated for non-inverted drones; blade drag is rotated from body to world coordinates.
  • Optional aerodynamic effects: For vertically ordered drone pairs separated laterally by less than 10 m, the lower drone receives a downward downwash force.The Crazyflie 2.x coefficients are (C1, C2, C3) = (2267.18, 0.16, −0.11).

3.2 Physics Modes

Section 3.2 presents six selectable physics modes, including MuJoCo-integrated rigid-body dynamics through the MJC* family and explicit Python Euler integration through DYN.

  • Physics Modes: Six selectable physics modes are listed in Table 1.
  • Physics Modes: The MJC* family applies rotor forces and torques through MuJoCo’s external-force interface, allowing mj step to integrate rigid-body dynamics.
  • Physics Modes: DYN explicitly integrates body dynamics in Python using Euler integration, writes the resulting pose into qpos / qvel, and then calls mj forward.

3.3 Drone Models

The simulator bundles three rigid-body drone models, parameterized for Crazyflies and RACE, with derived hover-speed and motor-saturation quantities computed during initialization.

  • Drone models: Three rigid-body models use Crazyflie parameters from and a published racing-class parameter set for RACE.These models are bundled with the environment and summarized in Table 2.
  • Derived quantities: BaseAviary computes HOVER RPM as the hover rotor speed and MAX RPM as the motor-saturation limit.Initialization uses m, g, kf, and T/W.
  • Downwash: Under Physics.MJC GND DRAG DW, the lower drone increases thrust to compensate for the upper drone’s wake.The behavior is illustrated with two stacked Crazyflies in examples/downwash.py.

3.4 Software Architecture and APIs

The package organizes task environments, controllers, wrappers, utilities, and a GPU vectorized backend into a modular Python architecture. Its environments share a Gymnasium-based API with configurable simulation, action, observation, and installation options.

  • Package organization: The envs/ module provides seven Gymnasium tasks plus a PettingZoo multi-agent wrapper.It includes base, hover, velocity, multi-hover, fly-through, formation, race, and multi-agent aviary modules.
  • Package organization: The control/ module implements paper-grade PIDControl and DSL-style DSLPIDControl with rate limiting and anti-windup.
  • Package organization: The wrappers/ module supplies wind, curriculum, and obstacle wrappers, while vectorized/ contains the MJX-based GPU backend.
  • Installation: The package is pip-installable with dependency groups for reinforcement learning, multi-agent learning, visualization, GPU acceleration, or all features.The groups add Stable-Baselines3, PettingZoo, matplotlib and Pillow, and jax[cuda12] with mujoco-mjx, respectively.
  • Environment API: All environments derive from BaseAviary(gym.Env) and use the standard Gymnasium five-value step interface.
  • Environment API: Constructor options configure drone models, fleet size, physics, frequencies, initial states, observations, actions, obstacles, vision, rendering, and output.sim freq must be a positive integer multiple of ctrl freq, producing sim freq/ctrl freq MuJoCo sub-steps per env.step call.

3.5 Composable Wrappers

Composable wrappers extend the simulator with wind, procedural obstacles, adaptive curricula, and per-episode domain randomization. These components can be attached during construction or reset-time configuration to vary disturbances, environments, task difficulty, and dynamics.

  • Wind: WindWrapper composes with any BaseAviary and applies configurable wind forces at every env.step, while BaseAviary supports in-place attachment.WindConfig supports CONSTANT, GUST, DRYDEN, SINUSOIDAL, and COMBINED models; COMBINED sums CONSTANT, GUST, DRYDEN, and SINUSOIDAL.
  • Procedural obstacles: ObstacleConfig selects procedural FOREST, URBAN, INDOOR, RANDOM, GATES, or CUSTOM generators and splices generated obstacle records into MuJoCo XML.A safe-zone radius keeps a sphere around each drone collision-free at spawn time.
  • Curriculum: CurriculumWrapper tracks a discrete difficulty level and updates it after episodes using moving success rate, reward, or episode length.A user-supplied difficulty function runs on every reset and can mutate the environment, such as raising target altitude or attaching wind as difficulty increases.
  • Domain randomization: DomainRandomizationWrapper resamples physical and sensor parameters at every reset while keeping them fixed within each episode, broadening training dynamics for sim-to-real transfer.The recommendations include randomizing coefficients rather than forces, including motor lag and action latency, adding sticky per-episode observation bias, and training wind through WindModel.COMBINED.

3.6 GPU-Vectorized Back End (MJX)

The MJXVectorAviary provides a fully GPU-vectorized, JAX-native back end that runs thousands of independent drone simulations in parallel and supports batched JAX workflows. Its XLA-vectorized execution scales to a default of 4096 environments, while a Gymnasium wrapper preserves compatibility with NumPy-based RL libraries.

  • GPU-vectorized execution: MJXVectorAviary runs thousands of independent drone simulations in parallel on a single GPU using MuJoCo MJX and JAX.The aviary can be batched with jax.vmap for population-based search and JAX-native pipelines such as PureJaxRL.
  • GPU-vectorized execution: A pure-JAX single-environment step is lifted across environments with jit(vmap(...)), allowing one XLA kernel to fan out across the batch.The batched state stores mjx.Data, a step counter, a PRNG key, and a done flag in an MJXState NamedTuple.
  • Scaling: The default configuration is num envs = 4096, with GPU memory as the only practical limit.The vectorized back end builds a minimal MJCF without visual meshes to keep XLA compilation tractable.
  • Interoperability: An MJXVecEnvGymWrapper bridges the JAX-native interface to Gymnasium VectorEnv by performing GPU →CPU transfers on each step().When the policy is written in JAX, the wrapper can be bypassed so the environment–policy–optimizer loop remains on device and avoids host / device traffic.
  • Task coverage: The MJX back end currently supports the hover, stabilize, and track tasks.This task support is stated as a current limitation of the back end.

4 Tasks, Control, and Examples

The section presents seven task environments, a modular control stack, and representative control and reinforcement-learning examples. It includes PettingZoo multi-agent integration and reports centimeter-scale tracking with an SE(3) geometric controller.

  • Task environments: The package ships seven task environments spanning hover, velocity tracking, multi-drone hover, waypoint navigation, formation flight, gate racing, and generic multi-agent use.The section organizes these environments in Section 4.1 and describes representative task objectives, rewards, and episode limits.
  • Multi-agent reinforcement learning: The PettingZoo ParallelEnv wrapper exposes per-agent 13-D observations, 4-D RPM actions, and per-drone hover rewards for multi-agent reinforcement-learning libraries.A 12-drone antipodal-navigation benchmark demonstrates collision avoidance through reward shaping or a safety-filter constraint under the same ParallelEnv contract.
  • Control stack: The control stack includes cascaded PID, DSLPIDControl with anti-windup and setpoint rate limits, and inline action handling for waypoint, velocity, or attitude setpoints.DSLPIDControl adds feed-forward acceleration and is preferred for aggressive trajectory tracking, while PIDControl converts desired thrust and torques into per-motor RPMs.
  • Control stack: 0.7 cm average tracking error and 3.5 cm peak error are achieved on a vertical circle with the SE(3) geometric controller in MuJoCo dynamics.The trajectory has radius 0.3 m and period 8 s; the peak occurs at the apex.
  • Examples and learning: The examples cover PID hover, DSLPID square-waypoint tracking, three-drone downwash control, PPO training with Stable-Baselines3, policy playback, and smoke tests across six environments.The examples use a 4-way SubprocVecEnv for PPO training, while the test suite covers all seven environments, controllers, and the PettingZoo contract.

5 Conclusions and Future Work

MuJoCo-Drones-Gym is presented as an open-source, Gymnasium- and PettingZoo-compatible multi-quadrotor environment built on MuJoCo, combining broad modularity with GPU-vectorized training support. Future work targets extending the MJX backend beyond rigid-body physics and expanding its observation and training capabilities.

  • Conclusions: MuJoCo-Drones-Gym is an open-source, Gymnasium- and PettingZoo-compatible multi-quadrotor environment built on MuJoCo.It preserves the API and research-oriented spirit of gym-pybullet-drones while using MuJoCo’s contact handling, integrator quality, and modern Python tooling.
  • Conclusions: The environment provides seven task templates, five action interfaces, three observation types, six physics modes, and three drone models.Physics options include toggleable ground effect, drag, and downwash.
  • Conclusions: Composable wind, obstacle, curriculum, and domain-randomization wrappers complement a GPU-vectorized MJX backend for batched JAX-native training.The wrappers and backend broaden the environment’s training and experiment-configuration capabilities.
  • Future Work: The MJXVectorAviary backend currently supports rigid-body MuJoCo physics only, leaving extensions for additional physics models, multi-drone observations, and vision-based observations.The paper identifies XLA-compilable ground-effect, drag, and downwash models as future work to enable full-task-suite PureJaxRL- and Brax-style pipelines.
  • Release: MuJoCo-Drones-Gym is released under the upstream repository’s license, with contributions, issues, and pull requests welcomed.The passage provides the project repository URL: https://github.com/tau-intelligence/MuJoCo-drones-gym.
Loading 2606.08039v1…