MuJoCo drone simulation with keyboard teleop

[Shreyas0812]

Problem

DimOS only supported drones over MAVLink / real hardware (dimos/robot/drone/). There was no way to bring up and control a drone in simulation, so drone control logic couldn't be developed or tested without a physical aircraft.

Closes #1842

Solution

Adds a MuJoCo based drone simulation with keyboard teleoperation:

• drone-basic-sim blueprint plus a drone_primitive_no_nav primitive and a DroneSimConnection (mujoco_sim.py) — no camera/lidar dependency.
• MuJoCo integration: loads the Crazyflie (cf2) model from MuJoCo Playground, spawns at z=0.5, and skips camera/lidar setup when the model doesn't provide them.
• DroneController (policy.py): maps keyboard input to desired tilt / yaw-rate / thrust — PD attitude control on roll and pitch, yaw-rate control with a max_yaw_rate limit, and thrust at hover with a lower floor (>=50% hover) for controlled slow descent.
• Keyboard teleop: added R/F altitude keys to KeyboardTeleop and routed tele_cmd_vel through to the sim connection.
• Viewer: MuJoCo camera now follows the drone.

How to Test

Requires the sim extra, which installs MuJoCo and MuJoCo Playground (source of the cf2 model):

uv sync --extra sim

dimos --simulation run drone-basic-sim

Drive with the keyboard (WASD / arrows for translation + yaw, R/F for altitude). The Crazyflie spawns in MuJoCo and is keyboard-controllable in real time.

Contributor License Agreement

• I have read and approved the CLA.

[greptile-apps]

Greptile Summary

Adds a MuJoCo-based Crazyflie (cf2) drone simulation with keyboard teleoperation, enabling drone control logic to be developed and tested without physical hardware.

• DroneController (policy.py) implements PD attitude control (roll/pitch) and yaw-rate control directly against data.ctrl[], running at the full simulation frequency (500 Hz for cf2) rather than the ONNX ctrl_dt cadence — appropriate for a direct-control drone.
• DroneMockController maps the four-dimensional drone command (forward, lateral, yaw, vertical) from shared memory and is selected when robot_model == "drone"; mujoco_process.py gains correct -1 guards for video and lidar rendering so models without those sensors skip the render path.
• KeyboardTeleop gains R/F and Up/Down arrow bindings for linear.z altitude control, with the multiplier, active-indicator, and help text all updated consistently.

Confidence Score: 5/5

The change is safe to merge — it adds a self-contained simulation path behind a new blueprint and does not touch any existing robot control flows.

All modifications to existing files are additive: new branches in match/if-elif blocks and new key bindings that don't alter existing behaviour. The drone simulation is gated behind its own blueprint and robot_model value, so other robots are unaffected.

No files require special attention — the core logic changes in mujoco_process.py are well-isolated to the cf2 robot branch.

Important Files Changed

Filename	Overview
dimos/simulation/mujoco/mujoco_process.py	Adds DroneMockController, drone branch in robot dispatch, spawn height for cf2, camera follow-lookat, and guards on camera/lidar rendering for models without those sensors.
dimos/simulation/mujoco/policy.py	Adds DroneController with PD attitude control, quaternion-to-Euler conversion, and direct ctrl[] writes; no substep throttling (runs at full sim freq), which is intentional for drone attitude loops.
dimos/simulation/mujoco/model.py	Adds cf2 assets loading, hover keyframe selection, and DroneController instantiation; DroneController receives ONNX params via **kwargs (intentional no-op for those fields).
dimos/robot/drone/mujoco_sim.py	New DroneSimConnection module wiring cmd_vel → MujocoConnection.move and odometry back to odom/tf streams; clean and minimal.
dimos/robot/unitree/keyboard_teleop.py	Adds R/F and Up/Down arrow bindings for linear.z (altitude); multiplier, active-check, indicator, and help text all updated consistently.
dimos/robot/drone/blueprints/basic/drone_basic_sim.py	New drone-basic-sim blueprint composing drone_primitive_no_nav, DroneSimConnection, and KeyboardTeleop; straightforward autoconnect composition.
dimos/robot/drone/blueprints/primitive/drone_primitive_no_nav.py	New no-nav drone primitive with vis/costmapper autoconnect, LCM transports for tele_cmd_vel/cmd_vel/odom, and basic Rerun visualization config.
dimos/robot/all_blueprints.py	Registers drone-basic-sim blueprint, drone-primitive-no-nav blueprint, and drone-sim-connection module; no issues.

Sequence Diagram

sequenceDiagram
    participant KT as KeyboardTeleop
    participant LCM as LCM /cmd_vel
    participant DSC as DroneSimConnection
    participant MC as MujocoConnection
    participant SHM as SharedMemory
    participant MP as mujoco_process
    participant DC as DroneController (cf2)

    KT->>LCM: tele_cmd_vel (Twist with linear.z)
    LCM->>DSC: cmd_vel
    DSC->>MC: move(twist)
    MC->>SHM: write_command(linear, angular)
    MP->>SHM: DroneMockController.get_command()
    MP->>DC: mj_step → set_mjcb_control callback
    DC->>DC: PD attitude control on roll/pitch/yaw/thrust
    DC-->>MP: data.ctrl[0..3] updated
    MP->>SHM: write_odom(pos, quat)
    SHM->>MC: odom_stream()
    MC->>DSC: _on_odom(SimOdometry)
    DSC->>DSC: odom.publish + tf.publish

Loading

_{Reviews (3): Last reviewed commit: "added up and down along with R/F keys" | Re-trigger Greptile}