dim-go2-sim 0.1.0__py3-none-any.whl

dim_go2_sim/__init__.py

@@ -0,0 +1,30 @@
+# Copyright 2026 Dimensional Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""A simulated Unitree Go2 for dimOS: cmd_vel in; Mid-360, camera, IMU, odometry, contacts out."""
+
+from dim_go2_sim.config import CameraConfig, LidarConfig, PhysicsConfig, SceneConfig
+from dim_go2_sim.contacts import Contact, Contacts
+from dim_go2_sim.module import Go2Sim, Go2SimConfig
+
+__all__ = [
+    "CameraConfig",
+    "Contact",
+    "Contacts",
+    "Go2Sim",
+    "Go2SimConfig",
+    "LidarConfig",
+    "PhysicsConfig",
+    "SceneConfig",
+]

dim_go2_sim/camera.py

@@ -0,0 +1,175 @@
+# Copyright 2026 Dimensional Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Camera frames drawn from a scene's textured glTF with pygfx (after PimSim's mesh renderer).
+
+Cooked scenes' MJCFs are collision-only, so MuJoCo would draw grey boxes. Only the static scene
+is drawn, not the robot: a camera must not see its own body. Rendering runs in a fresh child
+process, since GPU device state cannot cross dimOS's forkserver.
+"""
+
+from __future__ import annotations
+
+import json
+import math
+import pickle
+import subprocess
+import sys
+from typing import Any
+
+import numpy as np
+from numpy.typing import NDArray
+
+from dim_go2_sim.config import CameraConfig
+
+# glTF is Y-up; the world is Z-up.
+Y_UP_TO_Z_UP = np.array([[1, 0, 0, 0], [0, 0, -1, 0], [0, 1, 0, 0], [0, 0, 0, 1]], dtype=np.float64)
+
+
+def focal_length(config: CameraConfig) -> float:
+    return config.height / (2 * math.tan(math.radians(config.fovy_deg) / 2))
+
+
+class _Renderer:
+    def __init__(self, visual: str, y_up: bool, config: CameraConfig) -> None:
+        import pygfx as gfx  # type: ignore[import-untyped]
+        from rendercanvas.offscreen import OffscreenRenderCanvas  # type: ignore[import-untyped]
+
+        root = gfx.Group()
+        root.add(gfx.load_gltf(visual, quiet=True, remote_ok=False).scene)
+        if y_up:
+            root.local.matrix = Y_UP_TO_Z_UP
+        self._scene = gfx.Scene()
+        background = gfx.Background(None, gfx.BackgroundMaterial(config.background))
+        self._scene.add(background, _merged(gfx, root))
+        self._scene.add(gfx.AmbientLight("#ffffff", config.ambient))
+        sun = gfx.DirectionalLight("#fff4e6", config.sun)
+        sun.local.position = (4.0, -6.0, 10.0)
+        self._scene.add(sun)
+        canvas = OffscreenRenderCanvas(size=(config.width, config.height), pixel_ratio=1.0)
+        self._renderer = gfx.WgpuRenderer(canvas, pixel_ratio=1.0)
+        # pygfx's fov spans the mean of the view's height and width; match the vertical fovy.
+        aspect = config.width / config.height
+        fov = math.degrees(2 * math.atan(config.height * (1 + aspect) / (4 * focal_length(config))))
+        self._camera = gfx.PerspectiveCamera(
+            fov=fov, aspect=aspect, depth_range=(config.near, config.far)
+        )
+
+    def capture(
+        self, position: NDArray[np.float64], forward: NDArray[np.float64], up: NDArray[np.float64]
+    ) -> NDArray[np.uint8]:
+        self._camera.world.position = tuple(position)
+        self._camera.world.reference_up = tuple(up)
+        self._camera.look_at(tuple(position + forward))
+        self._renderer.render(self._scene, self._camera)
+        self._renderer.flush()
+        return np.ascontiguousarray(np.asarray(self._renderer.snapshot())[:, :, :3])
+
+
+def _merged(gfx: Any, root: Any) -> Any:
+    """One mesh per material, in world coordinates: a store is ~47k nodes, one draw call each."""
+    meshes: list[Any] = []
+    root.traverse(lambda node: meshes.append(node) if isinstance(node, gfx.Mesh) else None)
+    by_material: dict[int, list[Any]] = {}
+    for mesh in meshes:
+        by_material.setdefault(id(mesh.material), []).append(mesh)
+    group = gfx.Group()
+    for same in by_material.values():
+        positions, normals, texcoords, indices, offset = [], [], [], [], 0
+        for mesh in same:
+            geometry, matrix = mesh.geometry, mesh.world.matrix
+            local = geometry.positions.data.reshape(-1, 3).astype(np.float64)
+            positions.append((local @ matrix[:3, :3].T + matrix[:3, 3]).astype(np.float32))
+            normal = getattr(geometry, "normals", None)
+            if normal is not None:
+                rotated = normal.data.reshape(-1, 3) @ np.linalg.inv(matrix[:3, :3])
+                rotated /= np.maximum(np.linalg.norm(rotated, axis=1, keepdims=True), 1e-12)
+                normals.append(rotated.astype(np.float32))
+            uv = getattr(geometry, "texcoords", None)
+            if uv is not None:
+                texcoords.append(uv.data.reshape(-1, 2).astype(np.float32))
+            index = getattr(geometry, "indices", None)
+            index = np.arange(len(local)) if index is None else index.data.reshape(-1)
+            indices.append(index.astype(np.uint32) + offset)
+            offset += len(local)
+        attributes: dict[str, NDArray[Any]] = {
+            "positions": np.concatenate(positions),
+            "indices": np.concatenate(indices).reshape(-1, 3),
+        }
+        # A partial attribute can't be merged; drop it rather than misalign vertices.
+        if len(normals) == len(same):
+            attributes["normals"] = np.concatenate(normals)
+        if len(texcoords) == len(same):
+            attributes["texcoords"] = np.concatenate(texcoords)
+        group.add(gfx.Mesh(gfx.Geometry(**attributes), same[0].material))
+    return group
+
+
+def _serve(visual: str, y_up: bool, config: CameraConfig) -> None:
+    renderer = _Renderer(visual, y_up, config)
+    out = sys.stdout.buffer
+    pickle.dump(None, out)
+    out.flush()
+    while True:
+        try:
+            request = pickle.load(sys.stdin.buffer)
+        except EOFError:  # the parent exited without saying goodbye
+            break
+        if request is None:
+            break
+        pickle.dump(renderer.capture(*request), out, protocol=pickle.HIGHEST_PROTOCOL)
+        out.flush()
+
+
+class SceneCamera:
+    """Renders *visual* in a child process. Loading a large scene there takes a while."""
+
+    def __init__(self, visual: str, y_up: bool, config: CameraConfig) -> None:
+        self._process = subprocess.Popen(
+            [sys.executable, "-m", __name__, visual, json.dumps(y_up), config.model_dump_json()],
+            stdin=subprocess.PIPE,
+            stdout=subprocess.PIPE,
+        )
+        assert self._process.stdout is not None
+        try:
+            pickle.load(self._process.stdout)
+        except EOFError:
+            raise RuntimeError(
+                f"the camera renderer exited loading {visual}; see its stderr"
+            ) from None
+
+    def capture(
+        self, position: NDArray[np.float64], forward: NDArray[np.float64], up: NDArray[np.float64]
+    ) -> NDArray[np.uint8]:
+        assert self._process.stdin is not None and self._process.stdout is not None
+        pickle.dump((position, forward, up), self._process.stdin)
+        self._process.stdin.flush()
+        rgb: NDArray[np.uint8] = pickle.load(self._process.stdout)
+        return rgb
+
+    def close(self) -> None:
+        try:
+            assert self._process.stdin is not None
+            pickle.dump(None, self._process.stdin)
+            self._process.stdin.close()
+        except OSError:
+            pass
+        try:
+            self._process.wait(timeout=5)
+        except subprocess.TimeoutExpired:
+            self._process.kill()
+
+
+if __name__ == "__main__":
+    _serve(sys.argv[1], json.loads(sys.argv[2]), CameraConfig.model_validate_json(sys.argv[3]))

dim_go2_sim/config.py

@@ -0,0 +1,82 @@
+# Copyright 2026 Dimensional Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Everything about a Go2Sim run: the scene, the physics, and each sensor (no dimOS imports:
+the camera's child process loads this)."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from pydantic import BaseModel
+
+
+class SceneConfig(BaseModel):
+    # A cooked dimOS scene package by name (resolved with dimOS's catalog), or explicit files.
+    package: str | None = None
+    mjcf: Path | None = None  # collision geometry; a robot-free MJCF
+    visual: Path | None = None  # textured glTF/GLB the camera draws; None draws nothing
+    visual_y_up: bool = True  # glTF's convention; the world is Z-up
+    start: tuple[float, float, float] = (0.0, 0.0, 0.0)  # x, y, yaw of the robot
+    floor_z: float = 0.0  # height the feet start on
+    # Scene geometry whose top is below this counts as floor for contacts; the rest is clutter.
+    floor_top_z: float = 0.05
+
+
+class PhysicsConfig(BaseModel):
+    timestep: float = 0.0025  # seconds per MuJoCo step; the policy runs every 0.02 s
+    # The physics were fitted to real Go2 runs under the bundled policy (go2-sim, dimOS #4441).
+    fitted: bool = True
+    actuator_tau: float = 0.0151  # seconds of first-order lag between a torque and the joint
+    policy: Path | None = None  # an rl_sar-layout ONNX walking policy; None is the bundled one
+    # Trim the policy's command closed-loop on the true body velocity, as a Go2's sport mode
+    # does: the policy alone ignores slow commands and drifts when told to stand.
+    track_velocity: bool = True
+    command_timeout: float = 0.5  # seconds a cmd_vel holds before the robot stands still
+    max_velocity: tuple[float, float, float] = (1.0, 0.5, 1.5)  # vx, vy m/s and wz rad/s
+    realtime: bool = True  # pace to the wall clock; False runs as fast as it can
+    cache_dir: Path | None = None  # compiled worlds; None is the user cache directory
+
+
+class LidarConfig(BaseModel):
+    enabled: bool = True
+    rate_hz: float = 10.0
+    # The Mid-360 on a Go2's head, from the body origin, pitched nose-down.
+    xyz: tuple[float, float, float] = (0.29515, -0.00003, 0.16297)
+    pitch_deg: float = 60.0
+    point_rate_hz: int = 200_000
+    downsample: int = 1
+    min_range: float = 0.16
+    max_range: float = 40.0
+    noise: bool = True
+    dropout: bool = True
+    self_occlusion: bool = True  # the robot's own body blocks rays
+    motion_bins_hz: float = 200.0  # how often a scan's origin follows the moving body
+    seed: int = 0
+
+
+class CameraConfig(BaseModel):
+    enabled: bool = True
+    rate_hz: float = 10.0
+    width: int = 640
+    height: int = 480
+    fovy_deg: float = 60.0
+    # The Go2's nose camera, from the body origin, looking straight ahead.
+    xyz: tuple[float, float, float] = (0.32715, -0.00003, 0.04297)
+    pitch_deg: float = 0.0
+    near: float = 0.05
+    far: float = 100.0
+    ambient: float = 1.8
+    sun: float = 2.5
+    background: str = "#1b2230"

dim_go2_sim/contacts.py

@@ -0,0 +1,103 @@
+# Copyright 2026 Dimensional Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Which scene geometry the robot touches, by robot part and kind of geometry. Ground truth."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+import json
+import time
+from typing import Literal
+
+from dimos.types.timestamped import Timestamped
+from dimos_lcm.std_msgs import String as LCMString
+import mujoco
+import numpy as np
+
+Part = Literal["trunk", "leg", "foot"]
+Kind = Literal["floor", "clutter"]
+
+
+@dataclass(frozen=True, order=True)
+class Contact:
+    part: Part
+    kind: Kind
+
+
+class Contacts(Timestamped):
+    """The set of contacts at one instant, published whenever it changes."""
+
+    msg_name = "dim_go2_sim.Contacts"
+
+    def __init__(self, contacts: list[Contact] | None = None, ts: float | None = None) -> None:
+        self.contacts = contacts or []
+        self.ts = time.time() if ts is None else ts
+
+    def lcm_encode(self) -> bytes:
+        payload = {"ts": self.ts, "contacts": [[c.part, c.kind] for c in self.contacts]}
+        return LCMString(data=json.dumps(payload)).lcm_encode()  # type: ignore[no-any-return]
+
+    @classmethod
+    def lcm_decode(cls, data: bytes) -> Contacts:
+        payload = json.loads(LCMString.lcm_decode(data).data)
+        return cls([Contact(part, kind) for part, kind in payload["contacts"]], payload["ts"])
+
+
+class ContactSensor:
+    """Classifies each static scene geom once as floor (its top below floor_top_z) or clutter."""
+
+    def __init__(
+        self,
+        model: mujoco.MjModel,
+        data: mujoco.MjData,
+        robot_bodies: set[int],
+        trunk: int,
+        feet: tuple[int, ...],
+        floor_top_z: float,
+    ) -> None:
+        self._data = data
+        self._trunk = trunk
+        self._feet = set(feet)
+        self._robot = {g for g in range(model.ngeom) if int(model.geom_bodyid[g]) in robot_bodies}
+        self._body = model.geom_bodyid
+        self._kind: dict[int, Kind] = {}
+        for g in range(model.ngeom):
+            if g in self._robot or model.body_weldid[model.geom_bodyid[g]] != 0:
+                continue
+            if model.geom_type[g] == mujoco.mjtGeom.mjGEOM_PLANE:
+                self._kind[g] = "floor"
+                continue
+            up = data.geom_xmat[g].reshape(3, 3)[2]
+            center, half = model.geom_aabb[g, :3], model.geom_aabb[g, 3:]
+            top = data.geom_xpos[g, 2] + up @ center + np.abs(up) @ half
+            self._kind[g] = "floor" if top < floor_top_z else "clutter"
+
+    def contacts(self) -> list[Contact]:
+        """Every (robot part, scene kind) pair in contact right now, sorted."""
+        found: set[Contact] = set()
+        for i in range(self._data.ncon):
+            c = self._data.contact[i]
+            for scene, robot in ((c.geom1, c.geom2), (c.geom2, c.geom1)):
+                kind = self._kind.get(int(scene))
+                if kind is not None and int(robot) in self._robot:
+                    found.add(Contact(self._part(int(robot)), kind))
+        return sorted(found)
+
+    def _part(self, geom: int) -> Part:
+        if geom in self._feet:
+            return "foot"
+        if self._body[geom] == self._trunk:
+            return "trunk"
+        return "leg"

dim_go2_sim/data/go2/LICENSE

@@ -0,0 +1,27 @@
+Copyright (c) 2016-2022 HangZhou YuShu TECHNOLOGY CO.,LTD. ("Unitree Robotics")
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice, this
+  list of conditions and the following disclaimer.
+
+* Redistributions in binary form must reproduce the above copyright notice,
+  this list of conditions and the following disclaimer in the documentation
+  and/or other materials provided with the distribution.
+
+* Neither the name of the copyright holder nor the names of its
+  contributors may be used to endorse or promote products derived from
+  this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

dim_go2_sim/data/go2/README.md

@@ -0,0 +1,57 @@
+# Unitree Go2 Description (MJCF)
+
+> [!IMPORTANT]
+> Requires MuJoCo 2.2.2 or later.
+
+## Changelog
+
+See [CHANGELOG.md](./CHANGELOG.md) for a full history of changes.
+
+## Overview
+
+This package contains a simplified robot description (MJCF) of the [Go2
+Quadruped Robot](https://www.unitree.com/go2/) developed by [Unitree
+Robotics](https://www.unitree.com/). It is derived from the [publicly available
+URDF
+description](https://github.com/unitreerobotics/unitree_ros/tree/master/robots/go2_description).
+
+<p float="left">
+  <img src="go2.png" width="400">
+</p>
+
+## URDF → MJCF derivation steps
+
+1. Converted the DAE [mesh
+   files](https://github.com/unitreerobotics/unitree_ros/tree/master/robots/go2_description/meshes) to OBJ format using [Blender](https://www.blender.org/).
+   - When exporting, ensure "up axis" is `+Z`, and "forward axis" is `+Y`.
+2. Processed `.obj` files with [`obj2mjcf`](https://github.com/kevinzakka/obj2mjcf).
+3. Added `<mujoco> <compiler discardvisual="false" strippath="false" fusestatic="false"/> </mujoco>` to the URDF's
+   `<robot>` clause in order to preserve visual geometries.
+4. Loaded the URDF into MuJoCo and saved a corresponding MJCF.
+5. Added a `<freejoint/>` to the base.
+6. Manually edited the MJCF to extract common properties into the `<default>` section.
+7. Softened the contacts of the feet to approximate the effect of rubber and
+   increased `impratio` to reduce slippage.
+8. Added `scene.xml` which includes the robot, with a textured groundplane, skybox, and haze.
+
+## MJX
+
+A version of the Unitree Go2 model for use in MJX is available in `scene_mjx.xml` with the following changes:
+
+1. The solver `iterations` and `ls_iterations` were modified for performance.
+2. `frictionloss` was removed.
+3. The friction cone was changed from `elliptic` to `pyramidal`.
+4. All contacts other than the sphere geoms were changed to sphere geoms and placed at the joints positions to simulate collisions at critical locations.
+   - Note: Currently, MJX cannot deal with collision geoms other than spheres. In this case, we have to put many collision spheres in critical locations of the robot such as joints, torso, and feet as illustrated below.
+
+<p float="left">
+  <img src="go2_mjx.png" width="800">
+</p>
+
+## License
+
+This model is released under a [BSD-3-Clause License](LICENSE).
+
+## Acknowledgments
+
+The MJX model of Go2 was contributed by [lonelyfluency](https://github.com/lonelyfluency).