dm_control_env

Runtime error

App Files Files Community

Leeps commited on 6 days ago

Commit

dee27a2

1 Parent(s): 9b47d4a

Add drone + forest sim

Browse files

Files changed (5) hide show

Dockerfile +1 -0
server/__init__.py +2 -1
server/app.py +20 -2
server/drone_forest.xml +241 -0
server/drone_forest_environment.py +554 -0

Dockerfile CHANGED Viewed

@@ -71,4 +71,5 @@ HEALTHCHECK --interval=30s --timeout=3s --start-period=10s --retries=3 \
 # Run the FastAPI server
 # Use exec to replace the shell with uvicorn so it receives SIGINT/SIGTERM directly
 ENV ENABLE_WEB_INTERFACE=true
 CMD ["sh", "-c", "cd /app/env && exec uvicorn server.app:app --host 0.0.0.0 --port 8000"]

 # Run the FastAPI server
 # Use exec to replace the shell with uvicorn so it receives SIGINT/SIGTERM directly
 ENV ENABLE_WEB_INTERFACE=true
+ENV OPENENV_ENVIRONMENT="drone_forest"
 CMD ["sh", "-c", "cd /app/env && exec uvicorn server.app:app --host 0.0.0.0 --port 8000"]

server/__init__.py CHANGED Viewed

@@ -1,5 +1,6 @@
 """dm_control OpenEnv server module."""
 from .dm_control_environment import DMControlEnvironment
-__all__ = ["DMControlEnvironment"]

 """dm_control OpenEnv server module."""
 from .dm_control_environment import DMControlEnvironment
+from .drone_forest_environment import DroneForestEnvironment
+__all__ = ["DMControlEnvironment", "DroneForestEnvironment"]

server/app.py CHANGED Viewed

@@ -21,11 +21,14 @@ Usage:
     uv run --project . server
 """
 try:
     from openenv.core.env_server.http_server import create_app
     from ..models import DMControlAction, DMControlObservation
     from .dm_control_environment import DMControlEnvironment
 except ImportError:
     from openenv.core.env_server.http_server import create_app
@@ -38,25 +41,40 @@ except ImportError:
             sys.path.insert(0, _parent)
         from models import DMControlAction, DMControlObservation
         from server.dm_control_environment import DMControlEnvironment
     except ImportError:
         try:
             from dm_control_env.models import DMControlAction, DMControlObservation
             from dm_control_env.server.dm_control_environment import (
                 DMControlEnvironment,
             )
         except ImportError:
             from envs.dm_control_env.models import DMControlAction, DMControlObservation
             from envs.dm_control_env.server.dm_control_environment import (
                 DMControlEnvironment,
             )
 # Create the app with web interface
 # Pass the class (factory) for concurrent session support
 app = create_app(
-    DMControlEnvironment,
     DMControlAction,
     DMControlObservation,
-    env_name="dm_control_env",
 )

     uv run --project . server
 """
+import os
 try:
     from openenv.core.env_server.http_server import create_app
     from ..models import DMControlAction, DMControlObservation
     from .dm_control_environment import DMControlEnvironment
+    from .drone_forest_environment import DroneForestEnvironment
 except ImportError:
     from openenv.core.env_server.http_server import create_app
             sys.path.insert(0, _parent)
         from models import DMControlAction, DMControlObservation
         from server.dm_control_environment import DMControlEnvironment
+        from server.drone_forest_environment import DroneForestEnvironment
     except ImportError:
         try:
             from dm_control_env.models import DMControlAction, DMControlObservation
             from dm_control_env.server.dm_control_environment import (
                 DMControlEnvironment,
             )
+            from dm_control_env.server.drone_forest_environment import (
+                DroneForestEnvironment,
+            )
         except ImportError:
             from envs.dm_control_env.models import DMControlAction, DMControlObservation
             from envs.dm_control_env.server.dm_control_environment import (
                 DMControlEnvironment,
             )
+            from envs.dm_control_env.server.drone_forest_environment import (
+                DroneForestEnvironment,
+            )
+# Select environment based on OPENENV_ENVIRONMENT env var
+if os.environ.get("OPENENV_ENVIRONMENT") == "drone_forest":
+    _env_cls = DroneForestEnvironment
+    _env_name = "drone_forest"
+else:
+    _env_cls = DMControlEnvironment
+    _env_name = "dm_control_env"
 # Create the app with web interface
 # Pass the class (factory) for concurrent session support
 app = create_app(
+    _env_cls,
     DMControlAction,
     DMControlObservation,
+    env_name=_env_name,
 )

server/drone_forest.xml ADDED Viewed

	@@ -0,0 +1,241 @@

+<mujoco model="drone_forest">
+  <compiler angle="radian"/>
+  <option timestep="0.002" gravity="0 0 -9.81" integrator="RK4"/>
+  <default>
+    <geom contype="1" conaffinity="1" condim="3" friction="1 0.5 0.5"/>
+    <joint damping="0.01"/>
+  </default>
+  <asset>
+    <texture name="grid" type="2d" builtin="checker" rgb1="0.3 0.6 0.3" rgb2="0.2 0.5 0.2"
+             width="512" height="512"/>
+    <material name="ground_mat" texture="grid" texrepeat="10 10" texuniform="true"
+              reflectance="0.1"/>
+    <material name="drone_body_mat" rgba="0.2 0.2 0.2 1"/>
+    <material name="drone_arm_mat" rgba="0.4 0.4 0.4 1"/>
+    <material name="rotor_mat" rgba="0.1 0.1 0.8 1"/>
+    <material name="trunk_mat" rgba="0.45 0.3 0.15 1"/>
+    <material name="canopy_mat" rgba="0.1 0.5 0.1 0.8"/>
+    <material name="target_mat" rgba="1.0 0.0 0.0 0.6"/>
+  </asset>
+  <worldbody>
+    <!-- Ground plane -->
+    <geom name="ground" type="plane" size="25 25 0.1" material="ground_mat"
+          contype="1" conaffinity="1"/>
+    <!-- Drone -->
+    <body name="drone" pos="0 0 1.5">
+      <freejoint name="drone_joint"/>
+      <inertial pos="0 0 0" mass="0.4" diaginertia="0.003 0.003 0.005"/>
+      <!-- Central body -->
+      <geom name="drone_body" type="box" size="0.06 0.06 0.025" material="drone_body_mat"
+            mass="0.2" contype="1" conaffinity="1"/>
+      <!-- Arm +X (front-right) -->
+      <geom name="arm_fr" type="capsule" fromto="0.04 -0.04 0 0.14 -0.14 0"
+            size="0.01" material="drone_arm_mat" mass="0.02" contype="0" conaffinity="0"/>
+      <site name="rotor_fr" pos="0.14 -0.14 0.01" size="0.04 0.005" type="cylinder"
+            rgba="0.1 0.1 0.8 0.5"/>
+      <!-- Arm +X +Y (front-left) -->
+      <geom name="arm_fl" type="capsule" fromto="0.04 0.04 0 0.14 0.14 0"
+            size="0.01" material="drone_arm_mat" mass="0.02" contype="0" conaffinity="0"/>
+      <site name="rotor_fl" pos="0.14 0.14 0.01" size="0.04 0.005" type="cylinder"
+            rgba="0.1 0.1 0.8 0.5"/>
+      <!-- Arm -X -Y (back-right) -->
+      <geom name="arm_br" type="capsule" fromto="-0.04 -0.04 0 -0.14 -0.14 0"
+            size="0.01" material="drone_arm_mat" mass="0.02" contype="0" conaffinity="0"/>
+      <site name="rotor_br" pos="-0.14 -0.14 0.01" size="0.04 0.005" type="cylinder"
+            rgba="0.1 0.1 0.8 0.5"/>
+      <!-- Arm -X +Y (back-left) -->
+      <geom name="arm_bl" type="capsule" fromto="-0.04 0.04 0 -0.14 0.14 0"
+            size="0.01" material="drone_arm_mat" mass="0.02" contype="0" conaffinity="0"/>
+      <site name="rotor_bl" pos="-0.14 0.14 0.01" size="0.04 0.005" type="cylinder"
+            rgba="0.1 0.1 0.8 0.5"/>
+      <!-- FPV camera (front-facing, slightly down) -->
+      <camera name="drone_fpv" pos="0.08 0 -0.01" xyaxes="0 -1 0 0.1 0 1"
+              fovy="90"/>
+    </body>
+    <!-- Target (visual only) -->
+    <body name="target" pos="5 5 1.5">
+      <geom name="target_geom" type="sphere" size="0.3" material="target_mat"
+            contype="0" conaffinity="0"/>
+    </body>
+    <!-- 25 tree bodies - positions set at runtime via Python -->
+    <body name="tree_0" pos="3 2 0">
+      <geom name="trunk_0" type="cylinder" size="0.15 1.5" pos="0 0 1.5"
+            material="trunk_mat" contype="1" conaffinity="1"/>
+      <geom name="canopy_0" type="sphere" size="0.6" pos="0 0 3.2"
+            material="canopy_mat" contype="0" conaffinity="0"/>
+    </body>
+    <body name="tree_1" pos="5 -3 0">
+      <geom name="trunk_1" type="cylinder" size="0.15 1.5" pos="0 0 1.5"
+            material="trunk_mat" contype="1" conaffinity="1"/>
+      <geom name="canopy_1" type="sphere" size="0.6" pos="0 0 3.2"
+            material="canopy_mat" contype="0" conaffinity="0"/>
+    </body>
+    <body name="tree_2" pos="-4 4 0">
+      <geom name="trunk_2" type="cylinder" size="0.15 1.5" pos="0 0 1.5"
+            material="trunk_mat" contype="1" conaffinity="1"/>
+      <geom name="canopy_2" type="sphere" size="0.6" pos="0 0 3.2"
+            material="canopy_mat" contype="0" conaffinity="0"/>
+    </body>
+    <body name="tree_3" pos="7 1 0">
+      <geom name="trunk_3" type="cylinder" size="0.15 1.5" pos="0 0 1.5"
+            material="trunk_mat" contype="1" conaffinity="1"/>
+      <geom name="canopy_3" type="sphere" size="0.6" pos="0 0 3.2"
+            material="canopy_mat" contype="0" conaffinity="0"/>
+    </body>
+    <body name="tree_4" pos="-2 -5 0">
+      <geom name="trunk_4" type="cylinder" size="0.15 1.5" pos="0 0 1.5"
+            material="trunk_mat" contype="1" conaffinity="1"/>
+      <geom name="canopy_4" type="sphere" size="0.6" pos="0 0 3.2"
+            material="canopy_mat" contype="0" conaffinity="0"/>
+    </body>
+    <body name="tree_5" pos="1 6 0">
+      <geom name="trunk_5" type="cylinder" size="0.15 1.5" pos="0 0 1.5"
+            material="trunk_mat" contype="1" conaffinity="1"/>
+      <geom name="canopy_5" type="sphere" size="0.6" pos="0 0 3.2"
+            material="canopy_mat" contype="0" conaffinity="0"/>
+    </body>
+    <body name="tree_6" pos="-6 -1 0">
+      <geom name="trunk_6" type="cylinder" size="0.15 1.5" pos="0 0 1.5"
+            material="trunk_mat" contype="1" conaffinity="1"/>
+      <geom name="canopy_6" type="sphere" size="0.6" pos="0 0 3.2"
+            material="canopy_mat" contype="0" conaffinity="0"/>
+    </body>
+    <body name="tree_7" pos="4 -6 0">
+      <geom name="trunk_7" type="cylinder" size="0.15 1.5" pos="0 0 1.5"
+            material="trunk_mat" contype="1" conaffinity="1"/>
+      <geom name="canopy_7" type="sphere" size="0.6" pos="0 0 3.2"
+            material="canopy_mat" contype="0" conaffinity="0"/>
+    </body>
+    <body name="tree_8" pos="-3 3 0">
+      <geom name="trunk_8" type="cylinder" size="0.15 1.5" pos="0 0 1.5"
+            material="trunk_mat" contype="1" conaffinity="1"/>
+      <geom name="canopy_8" type="sphere" size="0.6" pos="0 0 3.2"
+            material="canopy_mat" contype="0" conaffinity="0"/>
+    </body>
+    <body name="tree_9" pos="6 4 0">
+      <geom name="trunk_9" type="cylinder" size="0.15 1.5" pos="0 0 1.5"
+            material="trunk_mat" contype="1" conaffinity="1"/>
+      <geom name="canopy_9" type="sphere" size="0.6" pos="0 0 3.2"
+            material="canopy_mat" contype="0" conaffinity="0"/>
+    </body>
+    <body name="tree_10" pos="-5 -4 0">
+      <geom name="trunk_10" type="cylinder" size="0.15 1.5" pos="0 0 1.5"
+            material="trunk_mat" contype="1" conaffinity="1"/>
+      <geom name="canopy_10" type="sphere" size="0.6" pos="0 0 3.2"
+            material="canopy_mat" contype="0" conaffinity="0"/>
+    </body>
+    <body name="tree_11" pos="2 -2 0">
+      <geom name="trunk_11" type="cylinder" size="0.15 1.5" pos="0 0 1.5"
+            material="trunk_mat" contype="1" conaffinity="1"/>
+      <geom name="canopy_11" type="sphere" size="0.6" pos="0 0 3.2"
+            material="canopy_mat" contype="0" conaffinity="0"/>
+    </body>
+    <body name="tree_12" pos="-1 7 0">
+      <geom name="trunk_12" type="cylinder" size="0.15 1.5" pos="0 0 1.5"
+            material="trunk_mat" contype="1" conaffinity="1"/>
+      <geom name="canopy_12" type="sphere" size="0.6" pos="0 0 3.2"
+            material="canopy_mat" contype="0" conaffinity="0"/>
+    </body>
+    <body name="tree_13" pos="8 -2 0">
+      <geom name="trunk_13" type="cylinder" size="0.15 1.5" pos="0 0 1.5"
+            material="trunk_mat" contype="1" conaffinity="1"/>
+      <geom name="canopy_13" type="sphere" size="0.6" pos="0 0 3.2"
+            material="canopy_mat" contype="0" conaffinity="0"/>
+    </body>
+    <body name="tree_14" pos="-7 5 0">
+      <geom name="trunk_14" type="cylinder" size="0.15 1.5" pos="0 0 1.5"
+            material="trunk_mat" contype="1" conaffinity="1"/>
+      <geom name="canopy_14" type="sphere" size="0.6" pos="0 0 3.2"
+            material="canopy_mat" contype="0" conaffinity="0"/>
+    </body>
+    <body name="tree_15" pos="3 7 0">
+      <geom name="trunk_15" type="cylinder" size="0.15 1.5" pos="0 0 1.5"
+            material="trunk_mat" contype="1" conaffinity="1"/>
+      <geom name="canopy_15" type="sphere" size="0.6" pos="0 0 3.2"
+            material="canopy_mat" contype="0" conaffinity="0"/>
+    </body>
+    <body name="tree_16" pos="-4 -7 0">
+      <geom name="trunk_16" type="cylinder" size="0.15 1.5" pos="0 0 1.5"
+            material="trunk_mat" contype="1" conaffinity="1"/>
+      <geom name="canopy_16" type="sphere" size="0.6" pos="0 0 3.2"
+            material="canopy_mat" contype="0" conaffinity="0"/>
+    </body>
+    <body name="tree_17" pos="7 -5 0">
+      <geom name="trunk_17" type="cylinder" size="0.15 1.5" pos="0 0 1.5"
+            material="trunk_mat" contype="1" conaffinity="1"/>
+      <geom name="canopy_17" type="sphere" size="0.6" pos="0 0 3.2"
+            material="canopy_mat" contype="0" conaffinity="0"/>
+    </body>
+    <body name="tree_18" pos="-8 0 0">
+      <geom name="trunk_18" type="cylinder" size="0.15 1.5" pos="0 0 1.5"
+            material="trunk_mat" contype="1" conaffinity="1"/>
+      <geom name="canopy_18" type="sphere" size="0.6" pos="0 0 3.2"
+            material="canopy_mat" contype="0" conaffinity="0"/>
+    </body>
+    <body name="tree_19" pos="0 -8 0">
+      <geom name="trunk_19" type="cylinder" size="0.15 1.5" pos="0 0 1.5"
+            material="trunk_mat" contype="1" conaffinity="1"/>
+      <geom name="canopy_19" type="sphere" size="0.6" pos="0 0 3.2"
+            material="canopy_mat" contype="0" conaffinity="0"/>
+    </body>
+    <body name="tree_20" pos="5 8 0">
+      <geom name="trunk_20" type="cylinder" size="0.15 1.5" pos="0 0 1.5"
+            material="trunk_mat" contype="1" conaffinity="1"/>
+      <geom name="canopy_20" type="sphere" size="0.6" pos="0 0 3.2"
+            material="canopy_mat" contype="0" conaffinity="0"/>
+    </body>
+    <body name="tree_21" pos="-6 6 0">
+      <geom name="trunk_21" type="cylinder" size="0.15 1.5" pos="0 0 1.5"
+            material="trunk_mat" contype="1" conaffinity="1"/>
+      <geom name="canopy_21" type="sphere" size="0.6" pos="0 0 3.2"
+            material="canopy_mat" contype="0" conaffinity="0"/>
+    </body>
+    <body name="tree_22" pos="8 3 0">
+      <geom name="trunk_22" type="cylinder" size="0.15 1.5" pos="0 0 1.5"
+            material="trunk_mat" contype="1" conaffinity="1"/>
+      <geom name="canopy_22" type="sphere" size="0.6" pos="0 0 3.2"
+            material="canopy_mat" contype="0" conaffinity="0"/>
+    </body>
+    <body name="tree_23" pos="-2 8 0">
+      <geom name="trunk_23" type="cylinder" size="0.15 1.5" pos="0 0 1.5"
+            material="trunk_mat" contype="1" conaffinity="1"/>
+      <geom name="canopy_23" type="sphere" size="0.6" pos="0 0 3.2"
+            material="canopy_mat" contype="0" conaffinity="0"/>
+    </body>
+    <body name="tree_24" pos="1 -4 0">
+      <geom name="trunk_24" type="cylinder" size="0.15 1.5" pos="0 0 1.5"
+            material="trunk_mat" contype="1" conaffinity="1"/>
+      <geom name="canopy_24" type="sphere" size="0.6" pos="0 0 3.2"
+            material="canopy_mat" contype="0" conaffinity="0"/>
+    </body>
+  </worldbody>
+  <!-- Site-based thrust actuators: force applied in site-local z (tilts with drone) -->
+  <actuator>
+    <general name="thrust_fr" site="rotor_fr" gear="0 0 1 0 0 0" ctrllimited="true"
+             ctrlrange="0 3" gainprm="1" biasprm="0 0 0"/>
+    <general name="thrust_fl" site="rotor_fl" gear="0 0 1 0 0 0" ctrllimited="true"
+             ctrlrange="0 3" gainprm="1" biasprm="0 0 0"/>
+    <general name="thrust_br" site="rotor_br" gear="0 0 1 0 0 0" ctrllimited="true"
+             ctrlrange="0 3" gainprm="1" biasprm="0 0 0"/>
+    <general name="thrust_bl" site="rotor_bl" gear="0 0 1 0 0 0" ctrllimited="true"
+             ctrlrange="0 3" gainprm="1" biasprm="0 0 0"/>
+  </actuator>
+  <!-- Overhead tracking camera -->
+  <camera name="tracking" pos="0 0 25" xyaxes="1 0 0 0 1 0" mode="trackcom"
+          fovy="60"/>
+</mujoco>

server/drone_forest_environment.py ADDED Viewed

	@@ -0,0 +1,554 @@

+"""
+Drone Forest Navigation Environment.
+A quadrotor drone navigates through a forest of columns (trees) to reach a target.
+The RL policy commands velocity (forward/left/up/turn) while a built-in PD flight
+controller handles low-level motor mixing.
+"""
+import base64
+import io
+import os
+import sys
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+from uuid import uuid4
+# Configure MuJoCo rendering backend before importing mujoco
+if "MUJOCO_GL" not in os.environ and sys.platform != "darwin":
+    os.environ.setdefault("MUJOCO_GL", "egl")
+import numpy as np
+try:
+    from openenv.core.env_server.interfaces import Environment
+    from ..models import DMControlAction, DMControlObservation, DMControlState
+except ImportError:
+    from openenv.core.env_server.interfaces import Environment
+    try:
+        import sys as _sys
+        from pathlib import Path as _Path
+        _parent = str(_Path(__file__).parent.parent)
+        if _parent not in _sys.path:
+            _sys.path.insert(0, _parent)
+        from models import DMControlAction, DMControlObservation, DMControlState
+    except ImportError:
+        try:
+            from dm_control_env.models import (
+                DMControlAction,
+                DMControlObservation,
+                DMControlState,
+            )
+        except ImportError:
+            from envs.dm_control_env.models import (
+                DMControlAction,
+                DMControlObservation,
+                DMControlState,
+            )
+# ---------------------------------------------------------------------------
+# Constants
+# ---------------------------------------------------------------------------
+NUM_TREES = 25
+ARENA_HALF = 10.0       # arena is 20x20 m
+MAX_ALTITUDE = 8.0
+MIN_ALTITUDE = 0.1
+TARGET_RADIUS = 0.5     # success if within this distance
+TREE_MIN_SPACING = 1.5  # min distance between tree centres
+SPAWN_CLEAR_RADIUS = 2.0  # keep trees away from spawn
+TARGET_MIN_DIST = 5.0   # target at least this far from spawn
+MAX_STEPS = 1000
+PHYSICS_DT = 0.002
+CONTROL_DT = 0.02       # 50 Hz control
+# Velocity limits
+MAX_XY_VEL = 3.0   # m/s
+MAX_Z_VEL = 2.0    # m/s
+MAX_YAW_RATE = 2.0 # rad/s
+# Flight-controller PD gains
+KP_VEL = 4.0
+KD_VEL = 1.5
+KP_ATT = 8.0
+KD_ATT = 2.0
+# Drone physical parameters
+DRONE_MASS = 0.48       # total mass (body 0.4 + arms 0.08) close to XML
+GRAVITY = 9.81
+HOVER_THRUST = DRONE_MASS * GRAVITY / 4.0  # per-motor hover
+ARM_LENGTH = 0.14       # distance from CoM to rotor
+XML_PATH = str(Path(__file__).parent / "drone_forest.xml")
+class DroneForestEnvironment(Environment):
+    """Drone navigates a randomised forest of columns to reach a target."""
+    SUPPORTS_CONCURRENT_SESSIONS = True
+    def __init__(
+        self,
+        render_height: Optional[int] = None,
+        render_width: Optional[int] = None,
+        **kwargs,
+    ):
+        self._model = None
+        self._data = None
+        self._render_height = render_height or int(
+            os.environ.get("DMCONTROL_RENDER_HEIGHT", "480")
+        )
+        self._render_width = render_width or int(
+            os.environ.get("DMCONTROL_RENDER_WIDTH", "640")
+        )
+        self._include_pixels = False
+        self._step_count = 0
+        self._prev_dist = None
+        self._tree_positions: List[np.ndarray] = []
+        self._target_pos = np.zeros(3)
+        self._done = False
+        self._rng = np.random.RandomState()
+        self._state = DMControlState(
+            episode_id=str(uuid4()),
+            step_count=0,
+            domain_name="drone_forest",
+            task_name="navigate",
+        )
+    # ------------------------------------------------------------------
+    # Model loading
+    # ------------------------------------------------------------------
+    def _ensure_model(self):
+        """Load MuJoCo model if not loaded yet."""
+        if self._model is not None:
+            return
+        import mujoco
+        self._model = mujoco.MjModel.from_xml_path(XML_PATH)
+        self._data = mujoco.MjData(self._model)
+        # Precompute body / geom ids
+        self._drone_body_id = mujoco.mj_name2id(
+            self._model, mujoco.mjtObj.mjOBJ_BODY, "drone"
+        )
+        self._target_body_id = mujoco.mj_name2id(
+            self._model, mujoco.mjtObj.mjOBJ_BODY, "target"
+        )
+        self._tree_body_ids = [
+            mujoco.mj_name2id(self._model, mujoco.mjtObj.mjOBJ_BODY, f"tree_{i}")
+            for i in range(NUM_TREES)
+        ]
+        self._trunk_geom_ids = [
+            mujoco.mj_name2id(self._model, mujoco.mjtObj.mjOBJ_GEOM, f"trunk_{i}")
+            for i in range(NUM_TREES)
+        ]
+        self._drone_body_geom_id = mujoco.mj_name2id(
+            self._model, mujoco.mjtObj.mjOBJ_GEOM, "drone_body"
+        )
+        self._ground_geom_id = mujoco.mj_name2id(
+            self._model, mujoco.mjtObj.mjOBJ_GEOM, "ground"
+        )
+        # Set state metadata
+        self._state.action_spec = {
+            "shape": [4],
+            "dtype": "float64",
+            "minimum": [-1.0, -1.0, -1.0, -1.0],
+            "maximum": [1.0, 1.0, 1.0, 1.0],
+            "name": "velocity_command",
+        }
+        self._state.observation_spec = {
+            "position": {"shape": [3], "dtype": "float64"},
+            "velocity": {"shape": [3], "dtype": "float64"},
+            "orientation": {"shape": [3], "dtype": "float64"},
+            "angular_velocity": {"shape": [3], "dtype": "float64"},
+            "target_relative": {"shape": [3], "dtype": "float64"},
+            "obstacle_distances": {"shape": [8], "dtype": "float64"},
+        }
+        self._state.physics_timestep = PHYSICS_DT
+        self._state.control_timestep = CONTROL_DT
+    # ------------------------------------------------------------------
+    # Forest randomisation
+    # ------------------------------------------------------------------
+    def _randomise_forest(self):
+        """Place trees and target using rejection sampling."""
+        import mujoco
+        positions = []
+        attempts = 0
+        while len(positions) < NUM_TREES and attempts < 5000:
+            x = self._rng.uniform(-ARENA_HALF + 1, ARENA_HALF - 1)
+            y = self._rng.uniform(-ARENA_HALF + 1, ARENA_HALF - 1)
+            # Keep clear of spawn
+            if np.sqrt(x ** 2 + y ** 2) < SPAWN_CLEAR_RADIUS:
+                attempts += 1
+                continue
+            # Min spacing from existing trees
+            ok = True
+            for p in positions:
+                if np.sqrt((x - p[0]) ** 2 + (y - p[1]) ** 2) < TREE_MIN_SPACING:
+                    ok = False
+                    break
+            if ok:
+                positions.append(np.array([x, y]))
+            attempts += 1
+        # Pad with far-away positions if we didn't get enough
+        while len(positions) < NUM_TREES:
+            positions.append(np.array([100.0, 100.0]))
+        self._tree_positions = positions
+        # Set tree body positions in the model
+        for i, pos in enumerate(positions):
+            body_id = self._tree_body_ids[i]
+            self._model.body_pos[body_id] = [pos[0], pos[1], 0.0]
+        # Place target: at least TARGET_MIN_DIST from origin, away from trees
+        for _ in range(1000):
+            angle = self._rng.uniform(0, 2 * np.pi)
+            dist = self._rng.uniform(TARGET_MIN_DIST, ARENA_HALF - 2)
+            tx, ty = dist * np.cos(angle), dist * np.sin(angle)
+            tz = self._rng.uniform(1.0, 3.0)
+            # Check clearance from trees
+            clear = True
+            for p in positions[:NUM_TREES]:
+                if np.sqrt((tx - p[0]) ** 2 + (ty - p[1]) ** 2) < 1.5:
+                    clear = False
+                    break
+            if clear:
+                break
+        self._target_pos = np.array([tx, ty, tz])
+        self._model.body_pos[self._target_body_id] = self._target_pos.copy()
+        # Recompute derived quantities after changing body positions
+        mujoco.mj_forward(self._model, self._data)
+    # ------------------------------------------------------------------
+    # Flight controller
+    # ------------------------------------------------------------------
+    def _flight_controller(self, cmd: np.ndarray) -> np.ndarray:
+        """
+        Convert velocity commands [vx, vy, vz, yaw_rate] in [-1,1]
+        to 4 motor thrusts.
+        """
+        # Scale commands
+        vx_cmd = cmd[0] * MAX_XY_VEL
+        vy_cmd = cmd[1] * MAX_XY_VEL
+        vz_cmd = cmd[2] * MAX_Z_VEL
+        yaw_rate_cmd = cmd[3] * MAX_YAW_RATE
+        # Current state
+        pos = self._data.qpos[:3].copy()
+        quat = self._data.qpos[3:7].copy()  # w, x, y, z
+        vel = self._data.qvel[:3].copy()
+        ang_vel = self._data.qvel[3:6].copy()
+        # Extract yaw from quaternion
+        roll, pitch, yaw = self._quat_to_euler(quat)
+        # Rotate desired world-frame velocity into body XY
+        cos_yaw, sin_yaw = np.cos(yaw), np.sin(yaw)
+        # World-frame desired velocity
+        vx_world = vx_cmd * cos_yaw - vy_cmd * sin_yaw
+        vy_world = vx_cmd * sin_yaw + vy_cmd * cos_yaw
+        # Velocity error in world frame
+        vx_err = vx_world - vel[0]
+        vy_err = vy_world - vel[1]
+        vz_err = vz_cmd - vel[2]
+        # Desired roll/pitch from XY velocity error (small angle approx)
+        desired_pitch = np.clip(KP_VEL * vx_err, -0.5, 0.5)
+        desired_roll = np.clip(-KP_VEL * vy_err, -0.5, 0.5)
+        # Attitude PD
+        roll_err = desired_roll - roll
+        pitch_err = desired_pitch - pitch
+        yaw_rate_err = yaw_rate_cmd - ang_vel[2]
+        torque_roll = KP_ATT * roll_err - KD_ATT * ang_vel[0]
+        torque_pitch = KP_ATT * pitch_err - KD_ATT * ang_vel[1]
+        torque_yaw = KP_ATT * yaw_rate_err
+        # Collective thrust: hover + vertical velocity correction
+        thrust = DRONE_MASS * GRAVITY + KP_VEL * vz_err * DRONE_MASS
+        # Quadrotor mixer: convert thrust + torques to 4 motor thrusts
+        # Layout: FR(+x,-y), FL(+x,+y), BR(-x,-y), BL(-x,+y)
+        L = ARM_LENGTH
+        t_fr = thrust / 4.0 + torque_pitch / (4.0 * L) - torque_roll / (4.0 * L) - torque_yaw / 4.0
+        t_fl = thrust / 4.0 + torque_pitch / (4.0 * L) + torque_roll / (4.0 * L) + torque_yaw / 4.0
+        t_br = thrust / 4.0 - torque_pitch / (4.0 * L) - torque_roll / (4.0 * L) + torque_yaw / 4.0
+        t_bl = thrust / 4.0 - torque_pitch / (4.0 * L) + torque_roll / (4.0 * L) - torque_yaw / 4.0
+        # Clamp to actuator range [0, 3]
+        motors = np.clip([t_fr, t_fl, t_br, t_bl], 0.0, 3.0)
+        return motors
+    @staticmethod
+    def _quat_to_euler(quat: np.ndarray):
+        """Convert quaternion [w, x, y, z] to Euler angles [roll, pitch, yaw]."""
+        w, x, y, z = quat
+        # Roll (x-axis rotation)
+        sinr = 2.0 * (w * x + y * z)
+        cosr = 1.0 - 2.0 * (x * x + y * y)
+        roll = np.arctan2(sinr, cosr)
+        # Pitch (y-axis rotation)
+        sinp = 2.0 * (w * y - z * x)
+        sinp = np.clip(sinp, -1.0, 1.0)
+        pitch = np.arcsin(sinp)
+        # Yaw (z-axis rotation)
+        siny = 2.0 * (w * z + x * y)
+        cosy = 1.0 - 2.0 * (y * y + z * z)
+        yaw = np.arctan2(siny, cosy)
+        return roll, pitch, yaw
+    # ------------------------------------------------------------------
+    # Observations
+    # ------------------------------------------------------------------
+    def _get_obs(self) -> Dict[str, List[float]]:
+        pos = self._data.qpos[:3].copy()
+        vel = self._data.qvel[:3].copy()
+        quat = self._data.qpos[3:7].copy()
+        ang_vel = self._data.qvel[3:6].copy()
+        roll, pitch, yaw = self._quat_to_euler(quat)
+        target_rel = self._target_pos - pos
+        # 8 nearest obstacle distances (XY plane, from drone position)
+        dists = []
+        for tp in self._tree_positions:
+            dx = tp[0] - pos[0]
+            dy = tp[1] - pos[1]
+            dists.append(np.sqrt(dx ** 2 + dy ** 2))
+        dists.sort()
+        obstacle_distances = dists[:8]
+        # Pad if fewer than 8
+        while len(obstacle_distances) < 8:
+            obstacle_distances.append(50.0)
+        return {
+            "position": pos.tolist(),
+            "velocity": vel.tolist(),
+            "orientation": [float(roll), float(pitch), float(yaw)],
+            "angular_velocity": ang_vel.tolist(),
+            "target_relative": target_rel.tolist(),
+            "obstacle_distances": obstacle_distances,
+        }
+    # ------------------------------------------------------------------
+    # Collision detection
+    # ------------------------------------------------------------------
+    def _check_collisions(self) -> bool:
+        """Return True if drone collides with any tree trunk or ground."""
+        import mujoco
+        for i in range(self._data.ncon):
+            contact = self._data.contact[i]
+            g1, g2 = contact.geom1, contact.geom2
+            pair = {g1, g2}
+            if self._drone_body_geom_id not in pair:
+                continue
+            other = (pair - {self._drone_body_geom_id}).pop()
+            if other == self._ground_geom_id or other in self._trunk_geom_ids:
+                return True
+        return False
+    # ------------------------------------------------------------------
+    # Reward
+    # ------------------------------------------------------------------
+    def _compute_reward(self, pos: np.ndarray) -> float:
+        dist = np.linalg.norm(self._target_pos - pos)
+        reward = 0.0
+        # Shaping: reward for getting closer
+        if self._prev_dist is not None:
+            reward += 1.0 * (self._prev_dist - dist)
+        self._prev_dist = dist
+        # Time pressure
+        reward -= 0.01
+        return float(reward)
+    # ------------------------------------------------------------------
+    # Termination
+    # ------------------------------------------------------------------
+    def _check_termination(self, pos: np.ndarray):
+        """Returns (done, bonus_reward)."""
+        dist = np.linalg.norm(self._target_pos - pos)
+        # Success
+        if dist < TARGET_RADIUS:
+            return True, 100.0
+        # Collision
+        if self._check_collisions():
+            return True, -50.0
+        # Out of bounds
+        if (abs(pos[0]) > ARENA_HALF or abs(pos[1]) > ARENA_HALF or
+                pos[2] > MAX_ALTITUDE or pos[2] < MIN_ALTITUDE):
+            return True, -10.0
+        # Max steps
+        if self._step_count >= MAX_STEPS:
+            return True, 0.0
+        return False, 0.0
+    # ------------------------------------------------------------------
+    # Core interface
+    # ------------------------------------------------------------------
+    def reset(
+        self,
+        domain_name: Optional[str] = None,
+        task_name: Optional[str] = None,
+        seed: Optional[int] = None,
+        render: bool = False,
+        **kwargs,
+    ) -> DMControlObservation:
+        import mujoco
+        self._ensure_model()
+        self._include_pixels = render
+        if seed is not None:
+            self._rng = np.random.RandomState(seed)
+        # Reset data to defaults
+        mujoco.mj_resetData(self._model, self._data)
+        # Randomise forest layout
+        self._randomise_forest()
+        # Place drone at origin, altitude 1.5
+        self._data.qpos[:3] = [0.0, 0.0, 1.5]
+        self._data.qpos[3:7] = [1.0, 0.0, 0.0, 0.0]  # identity quaternion
+        self._data.qvel[:] = 0.0
+        mujoco.mj_forward(self._model, self._data)
+        self._step_count = 0
+        pos = self._data.qpos[:3].copy()
+        self._prev_dist = float(np.linalg.norm(self._target_pos - pos))
+        self._done = False
+        self._state = DMControlState(
+            episode_id=str(uuid4()),
+            step_count=0,
+            domain_name="drone_forest",
+            task_name="navigate",
+            action_spec=self._state.action_spec,
+            observation_spec=self._state.observation_spec,
+            physics_timestep=PHYSICS_DT,
+            control_timestep=CONTROL_DT,
+        )
+        obs = self._get_obs()
+        pixels = self._render_pixels() if render else None
+        return DMControlObservation(
+            observations=obs,
+            pixels=pixels,
+            reward=0.0,
+            done=False,
+        )
+    def step(
+        self,
+        action: DMControlAction,
+        render: bool = False,
+        **kwargs,
+    ) -> DMControlObservation:
+        import mujoco
+        if self._model is None or self._data is None:
+            raise RuntimeError("Environment not initialized. Call reset() first.")
+        if self._done:
+            raise RuntimeError("Episode is done. Call reset() to start a new episode.")
+        # Clip action to [-1, 1]
+        cmd = np.clip(np.array(action.values[:4], dtype=np.float64), -1.0, 1.0)
+        # Run flight controller to get motor thrusts
+        motors = self._flight_controller(cmd)
+        # Set actuator controls
+        self._data.ctrl[:4] = motors
+        # Step physics for one control timestep (multiple physics substeps)
+        n_substeps = int(CONTROL_DT / PHYSICS_DT)
+        for _ in range(n_substeps):
+            mujoco.mj_step(self._model, self._data)
+        self._step_count += 1
+        self._state.step_count = self._step_count
+        pos = self._data.qpos[:3].copy()
+        # Compute reward and check termination
+        reward = self._compute_reward(pos)
+        done, bonus = self._check_termination(pos)
+        reward += bonus
+        self._done = done
+        obs = self._get_obs()
+        pixels = self._render_pixels() if (render or self._include_pixels) else None
+        return DMControlObservation(
+            observations=obs,
+            pixels=pixels,
+            reward=float(reward),
+            done=done,
+        )
+    async def reset_async(self, **kwargs) -> DMControlObservation:
+        if sys.platform == "darwin":
+            return self.reset(**kwargs)
+        else:
+            import asyncio
+            return await asyncio.to_thread(self.reset, **kwargs)
+    async def step_async(self, action: DMControlAction, render: bool = False, **kwargs) -> DMControlObservation:
+        if sys.platform == "darwin":
+            return self.step(action, render=render, **kwargs)
+        else:
+            import asyncio
+            return await asyncio.to_thread(self.step, action, render=render, **kwargs)
+    # ------------------------------------------------------------------
+    # Rendering
+    # ------------------------------------------------------------------
+    def _render_pixels(self) -> Optional[str]:
+        try:
+            import mujoco
+            renderer = mujoco.Renderer(self._model, height=self._render_height, width=self._render_width)
+            renderer.update_scene(self._data, camera="tracking")
+            frame = renderer.render()
+            renderer.close()
+            from PIL import Image
+            img = Image.fromarray(frame)
+            buf = io.BytesIO()
+            img.save(buf, format="PNG")
+            return base64.b64encode(buf.getvalue()).decode("utf-8")
+        except Exception:
+            return None
+    @property
+    def state(self) -> DMControlState:
+        return self._state
+    def close(self) -> None:
+        self._model = None
+        self._data = None
+    def __del__(self):
+        try:
+            self.close()
+        except Exception:
+            pass