From b25dc2d717cf98a30cb15c77cd94fb9021c0647c Mon Sep 17 00:00:00 2001 From: David Gable Date: Fri, 22 May 2026 08:54:20 -0700 Subject: [PATCH 1/8] Fixed dataframe namespace so you can chain methods without losing type hints --- src/mujoco_mojo/utils/dataframe.py | 43 +++++++++++++++++++++--------- 1 file changed, 31 insertions(+), 12 deletions(-) diff --git a/src/mujoco_mojo/utils/dataframe.py b/src/mujoco_mojo/utils/dataframe.py index 8f605133..8c790a8a 100644 --- a/src/mujoco_mojo/utils/dataframe.py +++ b/src/mujoco_mojo/utils/dataframe.py @@ -1,7 +1,10 @@ from __future__ import annotations from pathlib import Path -from typing import TYPE_CHECKING, Any, Protocol, TypedDict, cast +from typing import TYPE_CHECKING, Any, TypedDict, cast + +if TYPE_CHECKING: + from typing import Self import polars as pl from scipy.spatial.transform import Rotation as R @@ -28,13 +31,6 @@ logger = get_logger(__name__) -class MojoFrameProtocol(Protocol): - """Protocol to tell type checkers that .mojo is available on the DataFrame.""" - - @property - def mojo(self) -> MojoNamespace: ... - - class _MojoFrame(pl.DataFrame): """ Internal implementation of MojoFrame to house static loaders. @@ -78,10 +74,30 @@ def from_dict( if TYPE_CHECKING: # MojoFrame is seen by the IDE as the combination of: - # 1. Our loaders (_MojoFrame) + # 1. loaders (_MojoFrame) # 2. Polars methods (pl.DataFrame) - # 3. Our namespace (MojoFrameProtocol) - class MojoDataFrame(_MojoFrame, MojoFrameProtocol): ... + # 3. namespace (MojoFrameProtocol) + class MojoDataFrame(_MojoFrame): + @property + def mojo(self) -> MojoNamespace: ... + + # override polars methods that return DataFrame so the type propagates + def select(self, *args: Any, **kwargs: Any) -> Self: ... + def filter(self, *args: Any, **kwargs: Any) -> Self: ... + def with_columns(self, *args: Any, **kwargs: Any) -> Self: ... + def sort(self, *args: Any, **kwargs: Any) -> Self: ... + def head(self, *args: Any, **kwargs: Any) -> Self: ... + def tail(self, *args: Any, **kwargs: Any) -> Self: ... + def limit(self, *args: Any, **kwargs: Any) -> Self: ... + def slice(self, *args: Any, **kwargs: Any) -> Self: ... + def rename(self, *args: Any, **kwargs: Any) -> Self: ... + def drop(self, *args: Any, **kwargs: Any) -> Self: ... + def drop_nulls(self, *args: Any, **kwargs: Any) -> Self: ... + def unique(self, *args: Any, **kwargs: Any) -> Self: ... + def sample(self, *args: Any, **kwargs: Any) -> Self: ... + def join(self, *args: Any, **kwargs: Any) -> Self: ... + def hstack(self, *args: Any, **kwargs: Any) -> Self: ... + def vstack(self, *args: Any, **kwargs: Any) -> Self: ... else: # At runtime, it's just our internal class MojoDataFrame = _MojoFrame @@ -93,7 +109,6 @@ class ColumnManifest(TypedDict): available_quats: list[str] -@pl.api.register_dataframe_namespace("mojo") class MojoNamespace: """ Enhanced Polars DataFrame for MuJoCo Mojo telemetry. @@ -370,3 +385,7 @@ def with_filters( target_cols = columns or self._df.columns filter_map = {col: filters for col in target_cols} return self.with_filter_map(filter_map, omit_time=omit_time) + + +if not TYPE_CHECKING: + pl.api.register_dataframe_namespace("mojo")(MojoNamespace) From 8a1ac9f7f1d2dd89c6dd84627f101d811d212116 Mon Sep 17 00:00:00 2001 From: David Gable Date: Fri, 22 May 2026 09:47:24 -0700 Subject: [PATCH 2/8] Added better instructions for typings for MuJoCo --- docs/user-guides/getting-started.md | 8 ++++---- src/mujoco_mojo/__init__.py | 24 ++++++++++++++++++++++++ src/mujoco_mojo/utils/layers/cli.py | 19 +++++++++++++++++++ 3 files changed, 47 insertions(+), 4 deletions(-) diff --git a/docs/user-guides/getting-started.md b/docs/user-guides/getting-started.md index 1fb42c7a..4d7f42cd 100644 --- a/docs/user-guides/getting-started.md +++ b/docs/user-guides/getting-started.md @@ -283,23 +283,23 @@ MuJoCo does not currently ship first-party Python typing stubs. To enable proper > Thanks to the work by [@kevinzakka](https://github.com/google-deepmind/mujoco/issues/1292#issuecomment-1874138201) and [@mluogh-xdof](https://github.com/google-deepmind/mujoco/issues/1292#issuecomment-3208219200) for figuring all this out! -1. In your terminal: +1. From your project root, generate stubs into a local `typings/` directory: ```bash linenums="0" - pybind11-stubgen mujoco -o ~/typings/ --numpy-array-wrap-with-annotated + pybind11-stubgen mujoco -o typings/ --numpy-array-wrap-with-annotated ``` 2. Recent MuJoCo builds compiled with newer pybind11 versions correctly expose enums as `SupportsInt`. If you encounter Pyright enum type errors, apply the compatibility patch: ```bash linenums="0" - python ~/typings/patch_mujoco_enums.py ~/typings/mujoco/_enums.pyi + python typings/patch_mujoco_enums.py typings/mujoco/_enums.pyi ``` 3. Then in `pyproject.toml` (for me, VSCode already type hints correctly, but this should fix things if you use `pyright` with your pre-commit hooks): ```toml title="pyproject.toml" [tool.pyright] - stubPath = "~/typings" + stubPath = "typings" venvPath = "." venv = ".venv" ``` diff --git a/src/mujoco_mojo/__init__.py b/src/mujoco_mojo/__init__.py index c8cf36a6..2d029fd6 100644 --- a/src/mujoco_mojo/__init__.py +++ b/src/mujoco_mojo/__init__.py @@ -2,6 +2,30 @@ mujoco_mojo is a collection of Python objects built to make working with MuJoCo via Python easier. It provides vast bindings for all MJCF XML schema objects, tools to convert to XML, run MuJoCo simulations, and more. + +MuJoCo Type Stubs: +------------------ + +MuJoCo does not ship Python type stubs. Generate them once with pybind11-stubgen: + +.. code-block:: bash + pip install pybind11-stubgen + pybind11-stubgen mujoco -o typings/ --numpy-array-wrap-with-annotated + +Run this from the project root. The ``typings/`` directory is gitignored. +Then add to pyproject.toml: + +.. code-block:: toml + [tool.pyright] + stubPath = "typings" + venvPath = "." + venv = ".venv" + +If you encounter Pyright enum errors on recent MuJoCo builds, apply the +compatibility patch that ships alongside the stubs: + +.. code-block:: bash + python typings/patch_mujoco_enums.py typings/mujoco/_enums.pyi """ from mujoco_mojo.__about__ import __version__ # noqa: F401 diff --git a/src/mujoco_mojo/utils/layers/cli.py b/src/mujoco_mojo/utils/layers/cli.py index 48a049a0..988b1ef3 100644 --- a/src/mujoco_mojo/utils/layers/cli.py +++ b/src/mujoco_mojo/utils/layers/cli.py @@ -921,6 +921,25 @@ def init_project( border_style="cyan", ) ) + if not Path("typings", "mujoco").exists(): + console.print("\n[bold yellow]Type Hints Setup[/bold yellow]") + console.print( + "[bold white]MuJoCo does not ship Python type stubs.[/bold white] " + "Generate them once for Pylance/Pyright autocomplete:\n\n" + "[bold yellow]" + "pip install pybind11-stubgen\n" + "pybind11-stubgen mujoco -o typings/ --numpy-array-wrap-with-annotated" + "[/bold yellow]\n\n" + "Run from the project root. Then add to [bold cyan]pyproject.toml[/bold cyan]:\n\n" + "[dim]" + "\\[tool.pyright]\n" + 'stubPath = "typings"\n' + 'venvPath = "."\n' + 'venv = ".venv"' + "[/dim]\n\n" + "[dim]Enum errors? Run: " + "python typings/patch_mujoco_enums.py typings/mujoco/_enums.pyi[/dim]" + ) @cli_app.command(name="reloaded") From 1dd57fbc5e723a5ce481df5cdf155bfa8a25823e Mon Sep 17 00:00:00 2001 From: David Gable Date: Fri, 22 May 2026 10:20:57 -0700 Subject: [PATCH 3/8] Added fullscreen button to all pages for Dojo --- .../utils/layers/dojo/templates/base.html | 56 +- .../utils/layers/dojo/templates/morph.html | 52 +- .../partials/trial_viewer/_chart.html | 3 +- .../layers/dojo/templates/static/js/main.js | 10 +- .../dojo/templates/static/js/trial-viewer.js | 629 +++-- .../layers/dojo/templates/static/main.js | 295 --- .../layers/dojo/templates/static/monitor.js | 298 --- .../layers/dojo/templates/static/mosaic.js | 47 - .../dojo/templates/static/trial-viewer.js | 2032 ----------------- .../dojo/templates/static/ts/src/store.ts | 10 +- 10 files changed, 556 insertions(+), 2876 deletions(-) delete mode 100644 src/mujoco_mojo/utils/layers/dojo/templates/static/main.js delete mode 100644 src/mujoco_mojo/utils/layers/dojo/templates/static/monitor.js delete mode 100644 src/mujoco_mojo/utils/layers/dojo/templates/static/mosaic.js delete mode 100644 src/mujoco_mojo/utils/layers/dojo/templates/static/trial-viewer.js diff --git a/src/mujoco_mojo/utils/layers/dojo/templates/base.html b/src/mujoco_mojo/utils/layers/dojo/templates/base.html index d9917000..7095e6fa 100644 --- a/src/mujoco_mojo/utils/layers/dojo/templates/base.html +++ b/src/mujoco_mojo/utils/layers/dojo/templates/base.html @@ -139,9 +139,9 @@ -
-
+
+
@@ -194,9 +194,9 @@

- -
+ class="w-full md:w-auto flex flex-col gap-3 items-stretch"> + +
{% if status_code is not defined %} @@ -288,11 +288,40 @@

+ + + + + + + + {% endif %}

{% if status_code is not defined %} -
+
Owner

{{ owner }}

@@ -309,7 +338,8 @@

+ x-data + @keydown.escape.window="$store.dojo.exitFullscreen()">
+ :class="!$store.dojo.isPageReady ? 'opacity-0 blur-md scale-98 pointer-events-none' : 'opacity-100'"> {% block content %}{% endblock %}