Import upstream snapshot c19f713c415a699a79d71cd96aa13c3104a05047

Upstream: https://github.com/michaelgillett/mjlab
Upstream-Commit: c19f713c415a699a79d71cd96aa13c3104a05047
Upstream-Branch: main

docs/source/debugging/export_scene.rst

@@ -0,0 +1,50 @@
+.. _export-scene:
+
+Export Scene
+============
+
+The ``export-scene`` script writes a complete scene (XML and mesh assets) to a
+directory for inspection, sharing, or loading in standalone MuJoCo.
+
+Quick start
+-----------
+
+.. code-block:: bash
+
+    # Export a built-in entity by alias.
+    uv run export-scene g1 --output-dir /tmp/g1
+
+    # Export a registered task scene.
+    uv run export-scene Mjlab-Velocity-Flat-Unitree-Go1 --output-dir /tmp/task
+
+    # Export as a zip archive.
+    uv run export-scene yam --output-dir /tmp/yam --zip True
+
+    # Export a custom entity via import path.
+    uv run export-scene my_pkg.robots:get_my_robot_cfg --output-dir /tmp/custom
+
+The output directory contains a ``scene.xml`` and an ``assets/`` subdirectory
+with all referenced mesh files. The XML can be loaded directly with
+``mujoco.MjModel.from_xml_path()`` or dropped into the
+`simulate viewer <https://mujoco.readthedocs.io/en/stable/programming/samples.html#sasimulate>`_.
+
+Target resolution
+-----------------
+
+The positional ``target`` argument is resolved in order:
+
+1. **Task ID**: checked against the task registry (``import mjlab.tasks``).
+2. **Entity alias**: one of the built-in shorthands (``g1``, ``go1``, ``yam``).
+3. **Import path**: a ``module:attribute`` string pointing to any callable
+   that returns an ``EntityCfg``.
+
+If none match, the script prints available task IDs and aliases.
+
+Options
+-------
+
+``--output-dir DIR`` *(default: "export")*
+    Destination directory. Cleaned before each export to prevent stale assets.
+
+``--zip True`` *(default: False)*
+    Compress the output into a ``.zip`` archive and remove the directory.

docs/source/debugging/nan_guard.rst

@@ -0,0 +1,145 @@
+.. _nan-guard:
+
+NaN Guard
+=========
+
+The NaN guard captures simulation states when NaN/Inf is detected, helping
+debug numerical instability issues.
+
+
+Quick start
+-----------
+
+Enable the NaN guard with a single CLI flag:
+
+.. code-block:: bash
+
+    uv run train <task-name> --enable-nan-guard True
+
+This automatically captures and saves simulation states when NaN/Inf is
+detected. You can also enable it programmatically:
+
+.. code-block:: python
+
+    from mjlab.sim.sim import SimulationCfg
+    from mjlab.utils.nan_guard import NanGuardCfg
+
+    cfg = SimulationCfg(
+        nan_guard=NanGuardCfg(
+            enabled=True,
+            buffer_size=100,
+            output_dir="/tmp/mjlab/nan_dumps",
+            max_envs_to_dump=5,
+        ),
+    )
+
+
+Configuration
+-------------
+
+``enabled`` *(default: False)*
+    Enable/disable NaN detection and dumping.
+
+``buffer_size`` *(default: 100)*
+    Number of recent simulation states to keep in the rolling buffer.
+
+``output_dir`` *(default: "/tmp/mjlab/nan_dumps")*
+    Directory where NaN dump files are saved.
+
+``max_envs_to_dump`` *(default: 5)*
+    Maximum number of NaN environments to dump to disk. All environments are
+    tracked in the buffer, but only the first N are saved to reduce dump
+    size.
+
+
+Behavior
+--------
+
+- **Captures** simulation state before each step (``qpos``, ``qvel``,
+  ``act`` if the model has actuator activations, and ``mocap_pos``/``mocap_quat``
+  if the model has mocap bodies)
+- **Detects** NaN/Inf in ``qpos``, ``qvel``, ``qacc``,
+  ``qacc_warmstart``, and ``sensordata`` after each step
+- **Dumps** the rolling buffer and model to disk on first detection
+- **Stops** after the first dump to avoid spam
+
+When disabled, all operations are no-ops with negligible overhead.
+
+
+Output format
+-------------
+
+Each NaN detection creates timestamped files plus latest symlinks:
+
+- ``nan_dump_TIMESTAMP.npz``: compressed state buffer
+
+  - ``states_step_NNNNNN``: captured states per step
+    (shape: ``[num_envs_dumped, state_size]``)
+  - ``_metadata``: dict with ``num_envs_total``, ``nan_env_ids``,
+    ``dumped_env_ids``, etc.
+
+- ``model_TIMESTAMP.mjb``: MuJoCo model in binary format
+- ``nan_dump_latest.npz``: symlink to most recent dump
+- ``model_latest.mjb``: symlink to most recent model
+
+
+Visualizing dumps
+-----------------
+
+Use the interactive viewer to scrub through captured states:
+
+.. code-block:: bash
+
+    # View latest dump.
+    uv run viz-nan /tmp/mjlab/nan_dumps/nan_dump_latest.npz
+
+    # View a specific dump.
+    uv run viz-nan /tmp/mjlab/nan_dumps/nan_dump_20251014_123456.npz
+
+
+.. figure:: ../_static/content/nan_debug.gif
+   :alt: NaN Debug Viewer
+
+   NaN debug viewer.
+
+The viewer provides:
+
+- Step slider to scrub through the buffer
+- Environment slider to compare different environments
+- Info panel showing which environments have NaN/Inf
+- 3D visualization of the robot and terrain at each state
+
+
+NaN detection termination
+-------------------------
+
+While the NaN guard helps **debug** NaN issues by capturing states, you can
+also **prevent** training crashes using the ``nan_detection`` termination
+term. This marks NaN environments as terminated, allowing them to reset
+while training continues:
+
+.. code-block:: python
+
+    from mjlab.envs.mdp.terminations import nan_detection
+    from mjlab.managers.termination_manager import TerminationTermCfg
+
+    nan_term: TerminationTermCfg = field(
+        default_factory=lambda: TerminationTermCfg(
+            func=nan_detection,
+            time_out=False,
+        )
+    )
+
+Terminations are logged as ``Episode_Termination/nan_term`` in your metrics.
+
+.. important::
+
+   ``nan_detection`` is a band-aid, not a cure. If NaNs occur during your
+   task objective (e.g., NaNs happen when grasping), the policy will never
+   learn to complete the task since it resets before receiving rewards.
+   Monitor your ``Episode_Termination/nan_term`` metrics carefully.
+
+**When to use which:**
+
+- ``nan_guard``: debug and understand why NaNs occur (always do this first)
+- ``nan_detection``: keep training stable while working on a permanent fix