chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge remote-tracking branch 'origin/main' into tarik/mjx-warp-nested-dataclass-methods

Browse Source 
# Conflicts:
#	mjx/mujoco/mjx/warp/types.py
This commit is contained in:
Tarik Kelestemur
2026-06-03 11:36:53 -04:00
parent c80d8fb35d 0897e3aec1
commit 319b3975a0
68 changed files with 6108 additions and 1631 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/APIreference/functions.rst
+1 -44
View File
@@ -601,14 +601,7 @@ Get name of object with the specified :ref:`mjtObj` type and id, returns ``NULL`
.. mujoco-include:: mj_fullM
Convert sparse inertia matrix ``M`` into full (i.e. dense) matrix.
|br| ``dst`` must be of size ``nv x nv``, ``M`` must be of the same structure as ``mjData.qM``.
The ``mjData`` members ``qM`` and ``M`` represent the same matrix in different formats; the former is unique to
MuJoCo, the latter is standard Compressed Sparse Row (lower triangle only). The :math:`L^T D L` factor of the inertia
matrix ``mjData.qLD`` uses the same CSR format as ``mjData.M``. See
`engine_support_test <https://github.com/google-deepmind/mujoco/blob/main/test/engine/engine_support_test.cc>`__ for
pedagogical examples.
Convert sparse inertia matrix into full (i.e. dense) matrix.
.. _mj_mulM:
@@ -1980,24 +1973,6 @@ Error and memory
Main error function; does not return to caller.
.. _mju_error_i:
`mju_error_i <#mju_error_i>`__
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
.. mujoco-include:: mju_error_i
Deprecated: use mju_error.
.. _mju_error_s:
`mju_error_s <#mju_error_s>`__
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
.. mujoco-include:: mju_error_s
Deprecated: use mju_error.
.. _mju_warning:
`mju_warning <#mju_warning>`__
@@ -2007,24 +1982,6 @@ Deprecated: use mju_error.
Main warning function; returns to caller.
.. _mju_warning_i:
`mju_warning_i <#mju_warning_i>`__
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
.. mujoco-include:: mju_warning_i
Deprecated: use mju_warning.
.. _mju_warning_s:
`mju_warning_s <#mju_warning_s>`__
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
.. mujoco-include:: mju_warning_s
Deprecated: use mju_warning.
.. _mju_clearHandlers:
`mju_clearHandlers <#mju_clearHandlers>`__
doc/APIreference/functions_override.rst
-11
View File
@@ -342,17 +342,6 @@ found, the function will return ``distmax`` and ``fromto``, if given, will be se
As explained in :ref:`Collision Detection<coDistance>`, distances are inaccurate when using the
:ref:`legacy CCD pipeline<coCCD>`, and its use is discouraged.
.. _mj_fullM:
Convert sparse inertia matrix ``M`` into full (i.e. dense) matrix.
|br| ``dst`` must be of size ``nv x nv``, ``M`` must be of the same structure as ``mjData.qM``.
The ``mjData`` members ``qM`` and ``M`` represent the same matrix in different formats; the former is unique to
MuJoCo, the latter is standard Compressed Sparse Row (lower triangle only). The :math:`L^T D L` factor of the inertia
matrix ``mjData.qLD`` uses the same CSR format as ``mjData.M``. See
`engine_support_test <https://github.com/google-deepmind/mujoco/blob/main/test/engine/engine_support_test.cc>`__ for
pedagogical examples.
.. _mj_mulM:
This function multiplies the joint-space inertia matrix stored in ``mjData.M`` by a vector.
doc/XMLreference.rst
+1 -1
View File
@@ -3137,7 +3137,7 @@ Attributes may be applied or ignored depending on the lighting model being used.
.. _body-light-directional:
:at:`directional`: :at-val:`[false, true], "false"`
This is a deprecated legacy attribute. Please use :ref:`light <body-light-type>` type instead. If set to "true", and
This is a deprecated legacy attribute. Please use light :ref:`type <body-light-type>` instead. If set to "true", and
no type is specified, this will change the light type to be directional.
.. _body-light-castshadow:
doc/changelog.rst
+9 -5
View File
@@ -21,16 +21,20 @@ General
:class: attention
- The header file ``mjthread.h`` was removed along with the old engine threading API.
**Migration:** Use :ref:`mju_threadpool` to set number of worker threads for the engine.
|br| **Migration:** Use :ref:`mju_threadpool` to set number of worker threads for the engine.
- Moved island sparse matrix construction from :ref:`mj_island` (single threaded) into :ref:`mj_fwdConstraint`
(multi-threaded). The island-specific matrices ``iM, iLD, iefc_J`` were removed from the arena and are now
allocated on the stack.
- Following the introduction of the :ref:`diagexact<option-flag-diagexact>` flag, the ``mjData`` field
``efc_diagApprox`` was renamed to ``efc_diagA``, as it can now be either the exact or approximate diagonal of
the :math:`A` ("Delassus") matrix.
- The deprecated functions ``mju_{error,warning}_{i,s}`` have been removed.
- Changed the signature of :ref:`mj_fullM` from ``mj_fullM(m, dst, M)`` to ``mj_fullM(m, d, dst)`` as part of the
planned deprecation of ``mjData.qM`` in favor of the CSR-format ``mjData.M``.
**Migration:** For inertia matrix conversion, replace ``mj_fullM(m, dst, d->qM)`` with ``mj_fullM(m, d, dst)`` or
``mju_sym2dense(dst, d->M, m->nv, m->M_rownnz, m->M_rowadr, m->M_colind)``.
Bug fixes
^^^^^^^^^
@@ -2145,7 +2149,7 @@ General
<https://github.com/google-deepmind/mujoco/blob/main/model/humanoid/humanoid100.xml>`__ model, which previously
required ~500,000 ``mjtNum``'s, now only requires ~6000. Very large models can now load and run with the CG solver.
#. Modified :ref:`mju_error` and :ref:`mju_warning` to be variadic functions (support for printf-like arguments). The
functions :ref:`mju_error_i`, :ref:`mju_error_s`, :ref:`mju_warning_i`, and :ref:`mju_warning_s` are now deprecated.
functions ``mju_error_i``, ``mju_error_s``, ``mju_warning_i``, and ``mju_warning_s`` are now deprecated.
#. Implemented a performant ``mju_sqrMatTDSparse`` function that doesn't require dense memory allocation.
#. Added ``mj_stackAllocInt`` to get correct size for allocating ints on mjData stack. Reducing stack memory usage
by 10% - 15%.
doc/includes/references.h
+1 -5
View File
@@ -3320,7 +3320,7 @@ void mj_jacDot(const mjModel* m, const mjData* d, mjtNum* jacp, mjtNum* jacr,
void mj_angmomMat(const mjModel* m, mjData* d, mjtNum* mat, int body);
int mj_name2id(const mjModel* m, int type, const char* name);
const char* mj_id2name(const mjModel* m, int type, int id);
void mj_fullM(const mjModel* m, mjtNum* dst, const mjtNum* M);
void mj_fullM(const mjModel* m, const mjData* d, mjtNum* dst);
void mj_mulM(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec);
void mj_mulM2(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec);
void mj_addM(const mjModel* m, mjData* d, mjtNum* dst, int* rownnz, int* rowadr, int* colind);
@@ -3457,11 +3457,7 @@ void mjui_update(int section, int item, const mjUI* ui,
mjuiItem* mjui_event(mjUI* ui, mjuiState* state, const mjrContext* con);
void mjui_render(mjUI* ui, const mjuiState* state, const mjrContext* con);
void mju_error(const char* msg, ...) mjPRINTFLIKE(1, 2);
void mju_error_i(const char* msg, int i);
void mju_error_s(const char* msg, const char* text);
void mju_warning(const char* msg, ...) mjPRINTFLIKE(1, 2);
void mju_warning_i(const char* msg, int i);
void mju_warning_s(const char* msg, const char* text);
void mju_clearHandlers(void);
void* mju_malloc(size_t size);
void mju_free(void* ptr);
include/mujoco/experimental/usd/mjcPhysics/collisionAPI.h
+4
View File
@@ -323,6 +323,8 @@ class MjcPhysicsCollisionAPI : public UsdAPISchemaBase {
// --------------------------------------------------------------------- //
// MARGIN
// --------------------------------------------------------------------- //
/// DEPRECATED: Use newton:contactMargin and newton:contactGap instead.
///
/// Geometric inflation of the geom surface for the purpose of contact force
/// generation.
///
@@ -348,6 +350,8 @@ class MjcPhysicsCollisionAPI : public UsdAPISchemaBase {
// --------------------------------------------------------------------- //
// GAP
// --------------------------------------------------------------------- //
/// DEPRECATED: Use newton:contactGap instead.
///
/// Additional contact detection buffer beyond margin. Contacts are detected
/// at distance margin + gap but forces are only generated at distance margin.
///
include/mujoco/experimental/usd/mjcPhysics/equalityJointAPI.h
+82 -6
View File
@@ -43,12 +43,20 @@ class SdfAssetPath;
/// \class MjcPhysicsEqualityJointAPI
///
/// API providing extension attributes to represent equality/joint constraints.
/// This API is applied to a joint prim which acts as the constrained joint
/// (joint1 in MuJoCo terminology). The target relationship points to another
/// joint prim which is the reference joint (joint2 in MuJoCo terminology). The
/// constrained joint's position or angle is constrained to be a quartic
/// polynomial of the reference joint's position or angle. Only scalar joint
/// types (slide and hinge) can be used.
///
/// This API is applied to a joint prim which acts as the follower (joint0). The
/// leader joint (joint1) is specified via the newton:mimicJoint relationship
/// inherited from NewtonMimicAPI.
///
/// The follower's position or angle is constrained to be a quartic polynomial
/// of the leader's position or angle: joint0 = coef0 + coef1*(joint1) +
/// coef2*(joint1)^2 + coef3*(joint1)^3 + coef4*(joint1)^4
///
/// The constant (coef0) and linear (coef1) coefficients are provided by
/// NewtonMimicAPI as newton:mimicCoef0 and newton:mimicCoef1. The higher-order
/// coefficients (coef2-coef4) are provided by this API.
///
/// Only scalar joint types (slide and hinge) can be used.
///
class MjcPhysicsEqualityJointAPI : public UsdAPISchemaBase {
public:
@@ -150,10 +158,60 @@ class MjcPhysicsEqualityJointAPI : public UsdAPISchemaBase {
MJCPHYSICS_API
const TfType& _GetTfType() const override;
public:
// --------------------------------------------------------------------- //
// SOLIMP
// --------------------------------------------------------------------- //
/// Constraint solver parameter for equality constraint simulation.
///
/// | ||
/// | -- | -- |
/// | Declaration | `uniform double[] mjc:solimp = [0.9, 0.95, 0.001, 0.5, 2]`
/// | | C++ Type | VtArray<double> | | \ref Usd_Datatypes "Usd Type" |
/// SdfValueTypeNames->DoubleArray | | \ref SdfVariability "Variability" |
/// SdfVariabilityUniform |
MJCPHYSICS_API
UsdAttribute GetSolImpAttr() const;
/// See GetSolImpAttr(), and also
/// \ref Usd_Create_Or_Get_Property for when to use Get vs Create.
/// If specified, author \p defaultValue as the attribute's default,
/// sparsely (when it makes sense to do so) if \p writeSparsely is \c true -
/// the default for \p writeSparsely is \c false.
MJCPHYSICS_API
UsdAttribute CreateSolImpAttr(VtValue const& defaultValue = VtValue(),
bool writeSparsely = false) const;
public:
// --------------------------------------------------------------------- //
// SOLREF
// --------------------------------------------------------------------- //
/// Constraint solver parameter for equality constraint simulation.
///
/// | ||
/// | -- | -- |
/// | Declaration | `uniform double[] mjc:solref = [0.02, 1]` |
/// | C++ Type | VtArray<double> |
/// | \ref Usd_Datatypes "Usd Type" | SdfValueTypeNames->DoubleArray |
/// | \ref SdfVariability "Variability" | SdfVariabilityUniform |
MJCPHYSICS_API
UsdAttribute GetSolRefAttr() const;
/// See GetSolRefAttr(), and also
/// \ref Usd_Create_Or_Get_Property for when to use Get vs Create.
/// If specified, author \p defaultValue as the attribute's default,
/// sparsely (when it makes sense to do so) if \p writeSparsely is \c true -
/// the default for \p writeSparsely is \c false.
MJCPHYSICS_API
UsdAttribute CreateSolRefAttr(VtValue const& defaultValue = VtValue(),
bool writeSparsely = false) const;
public:
// --------------------------------------------------------------------- //
// COEF0
// --------------------------------------------------------------------- //
/// DEPRECATED: Use newton:mimicCoef0 instead.
///
/// Constant coefficient a0 of the quartic polynomial. The constraint is:
/// y = y0 + a0 + a1*(x-x0) + a2*(x-x0)^2 + a3*(x-x0)^3 + a4*(x-x0)^4.
///
@@ -179,6 +237,8 @@ class MjcPhysicsEqualityJointAPI : public UsdAPISchemaBase {
// --------------------------------------------------------------------- //
// COEF1
// --------------------------------------------------------------------- //
/// DEPRECATED: Use newton:mimicCoef1 instead.
///
/// Linear coefficient a1 of the quartic polynomial.
///
/// | ||
@@ -271,6 +331,22 @@ class MjcPhysicsEqualityJointAPI : public UsdAPISchemaBase {
UsdAttribute CreateCoef4Attr(VtValue const& defaultValue = VtValue(),
bool writeSparsely = false) const;
public:
// --------------------------------------------------------------------- //
// MJCTARGET
// --------------------------------------------------------------------- //
/// DEPRECATED: Use newton:mimicJoint instead.
///
/// Secondary target of the equality constraint (the leader/reference joint).
///
MJCPHYSICS_API
UsdRelationship GetMjcTargetRel() const;
/// See GetMjcTargetRel(), and also
/// \ref Usd_Create_Or_Get_Property for when to use Get vs Create
MJCPHYSICS_API
UsdRelationship CreateMjcTargetRel() const;
public:
// ===================================================================== //
// Feel free to add custom code below this line, it will be preserved by
include/mujoco/experimental/usd/mjcPhysics/materialAPI.h
+7
View File
@@ -42,6 +42,9 @@ class SdfAssetPath;
/// \class MjcPhysicsMaterialAPI
///
/// DEPRECATED: Use NewtonMaterialAPI instead. All attributes on this API have
/// been superseded by Newton equivalents.
///
/// API providing extension attributes to represent physical MuJoCo materials.
///
class MjcPhysicsMaterialAPI : public UsdAPISchemaBase {
@@ -148,6 +151,8 @@ class MjcPhysicsMaterialAPI : public UsdAPISchemaBase {
// --------------------------------------------------------------------- //
// TORSIONALFRICTION
// --------------------------------------------------------------------- //
/// DEPRECATED: Use newton:torsionalFriction instead.
///
/// Friction value acting around contact normal.
///
/// | ||
@@ -173,6 +178,8 @@ class MjcPhysicsMaterialAPI : public UsdAPISchemaBase {
// --------------------------------------------------------------------- //
// ROLLINGFRICTION
// --------------------------------------------------------------------- //
/// DEPRECATED: Use newton:rollingFriction instead.
///
/// Friction value acting around both axes on the contact tangent plane.
///
/// | ||
include/mujoco/experimental/usd/mjcPhysics/meshCollisionAPI.h
+2
View File
@@ -179,6 +179,8 @@ class MjcPhysicsMeshCollisionAPI : public UsdAPISchemaBase {
// --------------------------------------------------------------------- //
// MAXHULLVERT
// --------------------------------------------------------------------- //
/// DEPRECATED: Use newton:maxHullVertices instead.
///
/// Sets an upper limit on the number of vertices in the meshes convex hull.
/// The default value of -1 means unlimited.
///
include/mujoco/experimental/usd/mjcPhysics/sceneAPI.h
+6
View File
@@ -152,6 +152,8 @@ class MjcPhysicsSceneAPI : public UsdAPISchemaBase {
// --------------------------------------------------------------------- //
// TIMESTEP
// --------------------------------------------------------------------- //
/// DEPRECATED: Use newton:timeStepsPerSecond instead.
///
/// Controls the timestep in seconds used by MuJoCo.
///
/// | ||
@@ -498,6 +500,8 @@ class MjcPhysicsSceneAPI : public UsdAPISchemaBase {
// --------------------------------------------------------------------- //
// ITERATIONS
// --------------------------------------------------------------------- //
/// DEPRECATED: Use newton:maxSolverIterations instead.
///
/// Maximum number of iterations of the constraint solver.
///
/// | ||
@@ -942,6 +946,8 @@ class MjcPhysicsSceneAPI : public UsdAPISchemaBase {
// --------------------------------------------------------------------- //
// GRAVITYFLAG
// --------------------------------------------------------------------- //
/// DEPRECATED: Use newton:gravityEnabled instead.
///
/// Enables the application of gravitational acceleration as defined in
/// mjOption.
///
include/mujoco/experimental/usd/mjcPhysics/tokens.h
+3 -3
View File
@@ -656,7 +656,7 @@ struct MjcPhysicsTokensType {
const TfToken mjcSliderSite;
/// \brief "mjc:solimp"
///
/// MjcPhysicsCollisionAPI, MjcPhysicsEqualityAPI
/// MjcPhysicsCollisionAPI, MjcPhysicsEqualityAPI, MjcPhysicsEqualityJointAPI
const TfToken mjcSolimp;
/// \brief "mjc:solimpfriction"
///
@@ -672,7 +672,7 @@ struct MjcPhysicsTokensType {
const TfToken mjcSolmix;
/// \brief "mjc:solref"
///
/// MjcPhysicsCollisionAPI, MjcPhysicsEqualityAPI
/// MjcPhysicsCollisionAPI, MjcPhysicsEqualityAPI, MjcPhysicsEqualityJointAPI
const TfToken mjcSolref;
/// \brief "mjc:solreffriction"
///
@@ -700,7 +700,7 @@ struct MjcPhysicsTokensType {
const TfToken mjcStiffness;
/// \brief "mjc:target"
///
/// MjcPhysicsActuator, MjcPhysicsEqualityAPI
/// MjcPhysicsActuator, MjcPhysicsEqualityAPI, MjcPhysicsEqualityJointAPI
const TfToken mjcTarget;
/// \brief "mjc:torqueScale"
///
include/mujoco/mujoco.h
+2 -14
View File
@@ -598,8 +598,8 @@ MJAPI int mj_name2id(const mjModel* m, int type, const char* name);
// Get name of object with the specified mjtObj type and id; return NULL if name not found.
MJAPI const char* mj_id2name(const mjModel* m, int type, int id);
// Convert sparse inertia matrix M into full (i.e. dense) matrix.
MJAPI void mj_fullM(const mjModel* m, mjtNum* dst, const mjtNum* M);
// Convert sparse inertia matrix into full (i.e. dense) matrix.
MJAPI void mj_fullM(const mjModel* m, const mjData* d, mjtNum* dst);
// Multiply vector by inertia matrix.
MJAPI void mj_mulM(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec);
@@ -972,21 +972,9 @@ MJAPI void mjui_render(mjUI* ui, const mjuiState* state, const mjrContext* con);
// Main error function; does not return to caller.
MJAPI void mju_error(const char* msg, ...) mjPRINTFLIKE(1, 2);
// Deprecated: use mju_error.
MJAPI void mju_error_i(const char* msg, int i);
// Deprecated: use mju_error.
MJAPI void mju_error_s(const char* msg, const char* text);
// Main warning function; returns to caller.
MJAPI void mju_warning(const char* msg, ...) mjPRINTFLIKE(1, 2);
// Deprecated: use mju_warning.
MJAPI void mju_warning_i(const char* msg, int i);
// Deprecated: use mju_warning.
MJAPI void mju_warning_s(const char* msg, const char* text);
// Clear user error and memory handlers.
MJAPI void mju_clearHandlers(void);
mjx/mujoco/mjx/codegen/generate_warp_types.py
+15 -194
View File
@@ -17,23 +17,21 @@
import ast
import dataclasses
import enum
import inspect
import logging
import textwrap
import typing
from typing import Any, Callable, Dict, List, Optional, Set
from absl import app
from absl import flags
from etils import epath
import numpy as np
import warp as wp
import mujoco
from mujoco.mjx.codegen import file
import mujoco.mjx.third_party.mujoco_warp as mjwarp
import numpy as np
import warp as wp
from mujoco.mjx.third_party.warp._src.jax_experimental import ffi
_MJX_WARP_TYPES_OUT_FPATH = flags.DEFINE_string(
'mjx_warp_types_out_path',
'third_party/py/mujoco/mjx/warp/types.py',
@@ -47,7 +45,6 @@ _MJX_TYPES_PATH = flags.DEFINE_string(
)
_DATA_SHAPE_PROPERTY_FIELD = 'cacc'
_DOCSTRING_LINE_LENGTH = 80
_DUMMY_XML = """
<mujoco>
<worldbody>
@@ -64,29 +61,6 @@ _DUMMY_XML = """
"""
def _format_docstring(docstring: str) -> str:
"""Wraps docstring body lines to match checked-in generated files."""
formatted_lines = []
for line in docstring.splitlines():
if not line.strip():
formatted_lines.append(line)
continue
indent = line[: len(line) - len(line.lstrip())]
continuation_indent = indent + (' ' if indent else '')
formatted_lines.append(
textwrap.fill(
line,
width=_DOCSTRING_LINE_LENGTH,
subsequent_indent=continuation_indent,
break_long_words=False,
break_on_hyphens=False,
)
)
return '\n'.join(formatted_lines)
def _to_py_string(value, indent=0):
"""Converts a dictionary/set/tuple/type to a Python code string."""
indent_str = ' ' * indent
@@ -135,10 +109,8 @@ def _get_target_annotation_node(
if annotation == np.ndarray:
return _ast_parse_type('np.ndarray')
if (
isinstance(annotation, wp.array)
or type(annotation).__name__ == '_ArrayAnnotation'
):
if (isinstance(annotation, wp.array) or
type(annotation).__name__ == '_ArrayAnnotation'):
return _ast_parse_type('jax.Array')
if annotation in (int, float, bool):
@@ -194,10 +166,9 @@ def _get_annotations_recursive(
flattened = {}
for key, annotation in annotations.items():
full_key = f'{prefix}{key}'
if (
hasattr(annotation, '__annotations__')
and 'mujoco_warp' in annotation.__module__
):
if hasattr(
annotation, '__annotations__'
) and 'mujoco_warp' in annotation.__module__:
nested = _get_annotations_recursive(
dict(annotation.__annotations__), prefix=f'{full_key}__'
)
@@ -310,12 +281,6 @@ else:
Callback = None
PyTreeNode = mjx_dataclasses.PyTreeNode
def _as_numpy_array(value):
if hasattr(value, 'numpy'):
return value.numpy()
return np.asarray(value)
'''
target_fpath.write_text(header)
@@ -333,68 +298,6 @@ _FLATTEN_UNFLATTEN = """
"""
class _NumpyMethodAdapter(ast.NodeTransformer):
"""Adapts copied Warp array methods to MJX numpy-backed fields."""
def _is_asarray_call(self, node: ast.Call) -> bool:
return (
isinstance(node.func, ast.Attribute)
and isinstance(node.func.value, ast.Name)
and node.func.value.id == 'np'
and node.func.attr == 'asarray'
and len(node.args) == 1
and not node.keywords
)
def _is_as_numpy_array_call(self, node: ast.AST) -> bool:
return (
isinstance(node, ast.Call)
and isinstance(node.func, ast.Name)
and node.func.id == '_as_numpy_array'
)
def visit_Call(self, node: ast.Call) -> ast.AST: # pylint: disable=invalid-name
self.generic_visit(node)
if (
isinstance(node.func, ast.Attribute)
and node.func.attr == 'numpy'
and not node.args
and not node.keywords
):
return ast.copy_location(
ast.Call(
func=ast.Name(id='_as_numpy_array', ctx=ast.Load()),
args=[node.func.value],
keywords=[],
),
node,
)
if self._is_asarray_call(node) and self._is_as_numpy_array_call(
node.args[0]
):
return ast.copy_location(node.args[0], node)
return node
def _get_explicit_method_nodes(cls: Any) -> List[ast.FunctionDef]:
"""Returns explicit methods from a source class, adapted for MJX fields."""
source = textwrap.dedent(inspect.getsource(cls))
tree = ast.parse(source)
class_def = next(node for node in tree.body if isinstance(node, ast.ClassDef))
methods: List[ast.FunctionDef] = []
for node in class_def.body:
if not isinstance(node, ast.FunctionDef):
continue
if node.name in {'tree_flatten', 'tree_unflatten'}:
continue
methods.append(
typing.cast(ast.FunctionDef, _NumpyMethodAdapter().visit(node))
)
return methods
def write_nested_dataclass(target_fpath: epath.Path, cls: Any):
new_class_body = _build_new_class_body_ast(
set(cls.__annotations__.keys()),
@@ -402,89 +305,19 @@ def write_nested_dataclass(target_fpath: epath.Path, cls: Any):
dict(cls.__annotations__),
add_docstring=False,
)
new_class_body.extend(_get_explicit_method_nodes(cls))
cls_str = '\n'.join(
textwrap.indent(ast.unparse(node), ' ') for node in new_class_body
)
cls_str = '\n'.join([' ' + ast.unparse(node) for node in new_class_body])
cls_str = cls_str.replace('jax.Array', 'np.ndarray')
docstring = _format_docstring(cls.__doc__ or '')
with target_fpath.open('a') as f:
f.write(f'''
@dataclasses.dataclass(frozen=True)
@tree_util.register_pytree_node_class
class {cls.__name__}:
"""{docstring}"""
"""{cls.__doc__}"""
{cls_str}
{_FLATTEN_UNFLATTEN}
''')
def _get_class_name(line: str) -> str | None:
"""Returns the class name from a generated class definition line."""
if not line.startswith('class '):
return None
return line.removeprefix('class ').split('(', maxsplit=1)[0]
def _get_compact_derived_docstring_classes(
target_fpath: epath.Path,
) -> Set[str]:
"""Finds generated classes with no blank line after their docstring."""
compact_classes = set()
if not target_fpath.exists():
return compact_classes
lines = target_fpath.read_text().splitlines()
for i, line in enumerate(lines[:-2]):
cls_name = _get_class_name(line)
if cls_name is None:
continue
if (
'(PyTreeNode)' in line
and lines[i + 1].startswith(' """Derived fields from ')
and lines[i + 1].rstrip().endswith('."""')
and lines[i + 2].startswith(' ')
and ':' in lines[i + 2]
):
compact_classes.add(cls_name)
return compact_classes
def _restore_compact_derived_docstring_spacing(
target_fpath: epath.Path,
compact_classes: Set[str],
):
"""Restores compact generated PyTreeNode field class spacing."""
if not compact_classes:
return
lines = target_fpath.read_text().splitlines(keepends=True)
result = []
i = 0
current_class = None
while i < len(lines):
cls_name = _get_class_name(lines[i])
if cls_name is not None:
current_class = cls_name
result.append(lines[i])
if (
current_class in compact_classes
and lines[i].startswith(' """Derived fields from ')
and lines[i].rstrip().endswith('."""')
and i + 2 < len(lines)
and lines[i + 1].strip() == ''
and lines[i + 2].startswith(' ')
and ':' in lines[i + 2]
):
i += 2
continue
i += 1
target_fpath.write_text(''.join(result))
def _get_meta_fields(cls_name: str) -> Set[str]:
"""Returns the set of fields that should be meta-fields in the pytree."""
m = mujoco.MjModel.from_xml_string(_DUMMY_XML)
@@ -597,10 +430,8 @@ def _get_fields_with_cond(
if f.type in (int, float, bool) and add_static:
s.add(prefix + f.name)
continue
if not (
isinstance(f.type, wp.array)
or type(f.type).__name__ == '_ArrayAnnotation'
):
if not (isinstance(f.type, wp.array) or
type(f.type).__name__ == '_ArrayAnnotation'):
continue
if cond_fn(attr):
s.add(prefix + f.name)
@@ -639,10 +470,8 @@ batching.register_vmappable(DataWarp, int, int, _to_elt, _from_elt, None)
def _is_ffi_compatible(wp_type: Any) -> bool:
"""Returns True if the type is an array, scalar, or variadic tuple."""
if (
isinstance(wp_type, wp.array)
or type(wp_type).__name__ == '_ArrayAnnotation'
):
if (isinstance(wp_type, wp.array) or
type(wp_type).__name__ == '_ArrayAnnotation'):
return True
if wp_type in wp._src.types.value_types:
return True
@@ -719,9 +548,6 @@ def main(argv):
base_path = file.get_base_path()
target_fpath = base_path / _MJX_WARP_TYPES_OUT_FPATH.value
mjx_types_fpath = base_path / _MJX_TYPES_PATH.value
compact_derived_docstring_classes = _get_compact_derived_docstring_classes(
target_fpath
)
write_header(target_fpath)
# TODO(btaba): consider automated grabbing of nested dataclasses from mjwarp.
@@ -730,9 +556,7 @@ def main(argv):
write_core_cls('Statistic', target_fpath, mjx_types_fpath, set_diff=False)
write_core_cls(
'Option',
target_fpath,
mjx_types_fpath,
'Option', target_fpath, mjx_types_fpath,
extra_annotations={'graph_mode': ffi.GraphMode},
)
write_core_cls('Model', target_fpath, mjx_types_fpath)
@@ -743,9 +567,6 @@ def main(argv):
file.write_license(target_fpath)
file.format_file(target_fpath)
_restore_compact_derived_docstring_spacing(
target_fpath, compact_derived_docstring_classes
)
if __name__ == '__main__':
mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py
Vendored
+10 -2
View File
@@ -3058,11 +3058,11 @@ def create_render_context(
# Locate skybox texture
skybox_tex_ids = np.nonzero(mjm.tex_type == mujoco.mjtTexture.mjTEXTURE_SKYBOX)[0] if mjm.ntex else np.array([], dtype=int)
if render_skybox:
assert skybox_tex_ids.size > 0, "render_skybox=True but the model has no texture with type mjTEXTURE_SKYBOX"
if render_skybox and skybox_tex_ids.size > 0:
skybox_tex_id = int(skybox_tex_ids[0])
skybox_face_width = int(mjm.tex_width[skybox_tex_id])
else:
render_skybox = False
skybox_tex_id = -1
skybox_face_width = 1
@@ -3157,6 +3157,13 @@ def create_render_context(
bvh_ngeom = len(geom_enabled_idx)
# Geom types present among enabled geoms, plus FLEX when flex primitives exist.
# Used to statically eliminate unused intersection branches in the ray-cast kernels.
geom_ray_types = set(int(t) for t in mjm.geom_type[geom_enabled_idx])
if len(flex_geom_flexid) > 0:
geom_ray_types.add(int(types.GeomType.FLEX))
geom_ray_types = tuple(sorted(geom_ray_types))
rc = types.RenderContext(
nrender=ncam,
cam_res=cam_res_arr,
@@ -3210,6 +3217,7 @@ def create_render_context(
znear=znear,
total_rays=int(total),
enable_backface_culling=enable_backface_culling,
geom_ray_types=geom_ray_types,
)
bvh.build_scene_bvh(mjm, mjd, rc, nworld)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/render.py
Vendored
+338 -432
View File
@@ -151,449 +151,346 @@ def sample_skybox(
return wp.vec3(color[0], color[1], color[2])
# TODO: Investigate combining cast_ray and cast_ray_first_hit
@wp.func
def cast_ray(
# Model:
geom_type: wp.array[int],
geom_dataid: wp.array2d[int],
geom_size: wp.array2d[wp.vec3],
flex_vertadr: wp.array[int],
flex_edge: wp.array[wp.vec2i],
flex_radius: wp.array[float],
# Data in:
geom_xpos_in: wp.array2d[wp.vec3],
geom_xmat_in: wp.array2d[wp.mat33],
flexvert_xpos_in: wp.array2d[wp.vec3],
# In:
bvh_id: wp.uint64,
group_root: int,
worldid: int,
bvh_ngeom: int,
flex_bvh_ngeom: int,
enabled_geom_ids: wp.array[int],
mesh_bvh_id: wp.array[wp.uint64],
hfield_bvh_id: wp.array[wp.uint64],
flex_geom_flexid: wp.array[int],
flex_geom_edgeid: wp.array[int],
flex_bvh_id: wp.array[wp.uint64],
flex_group_root: wp.array2d[int],
ray_origin_world: wp.vec3,
ray_dir_world: wp.vec3,
cull_backfaces: bool,
) -> Tuple[int, float, wp.vec3, float, float, int, int]:
dist = float(MJ_MAXVAL)
normal = wp.vec3(0.0, 0.0, 0.0)
geom_id = int(-1)
bary_u = float(0.0)
bary_v = float(0.0)
face_idx = int(-1)
geom_mesh_id = int(-1)
def _make_cast_ray(geom_ray_types: Tuple[int], first_hit: bool = False) -> wp.Function:
"""Build a ray-cast func specialized to the geom types present in the scene.
query = wp.bvh_query_ray(bvh_id, ray_origin_world, ray_dir_world, group_root)
bounds_nr = int(0)
ngeom = bvh_ngeom + flex_bvh_ngeom
geom_ray_types is the set of GeomType int values that actually occur, so the
per-type intersection branches for absent types are eliminated at compile time
via wp.static, avoiding the register pressure of unreachable code paths.
while wp.bvh_query_next(query, bounds_nr, dist):
gi_global = bounds_nr
local_id = gi_global - (worldid * ngeom)
first_hit selects the variant (also resolved at compile time via wp.static):
- False: full closest-hit cast. Returns the closest hit's full surface data.
- True: any-hit cast (shadow rays). Uses the cheaper any-hit mesh/flex
intersections and returns on the first hit within max_dist. The result is
still the full tuple; callers test geom_id != -1 to detect a hit.
"""
d = float(-1.0)
hit_mesh_id = int(-1)
u = float(0.0)
v = float(0.0)
f = int(-1)
n = wp.vec3(0.0, 0.0, 0.0)
hit_geom_id = int(-1)
@wp.func
def cast_ray(
# Model:
geom_type: wp.array[int],
geom_dataid: wp.array2d[int],
geom_size: wp.array2d[wp.vec3],
flex_vertadr: wp.array[int],
flex_edge: wp.array[wp.vec2i],
flex_radius: wp.array[float],
# Data in:
geom_xpos_in: wp.array2d[wp.vec3],
geom_xmat_in: wp.array2d[wp.mat33],
flexvert_xpos_in: wp.array2d[wp.vec3],
# In:
bvh_id: wp.uint64,
group_root: int,
worldid: int,
bvh_ngeom: int,
flex_bvh_ngeom: int,
enabled_geom_ids: wp.array[int],
mesh_bvh_id: wp.array[wp.uint64],
hfield_bvh_id: wp.array[wp.uint64],
flex_geom_flexid: wp.array[int],
flex_geom_edgeid: wp.array[int],
flex_bvh_id: wp.array[wp.uint64],
flex_group_root: wp.array2d[int],
ray_origin_world: wp.vec3,
ray_dir_world: wp.vec3,
max_dist: float,
cull_backfaces: bool,
) -> Tuple[int, float, wp.vec3, float, float, int, int]:
dist = max_dist
normal = wp.vec3(0.0, 0.0, 0.0)
geom_id = int(-1)
bary_u = float(0.0)
bary_v = float(0.0)
face_idx = int(-1)
geom_mesh_id = int(-1)
if local_id < bvh_ngeom:
gi = enabled_geom_ids[local_id]
gtype = geom_type[gi]
else:
gi = local_id - bvh_ngeom
gtype = GeomType.FLEX
query = wp.bvh_query_ray(bvh_id, ray_origin_world, ray_dir_world, group_root)
bounds_nr = int(0)
ngeom = bvh_ngeom + flex_bvh_ngeom
hit_geom_id = gi
while wp.bvh_query_next(query, bounds_nr, dist):
gi_global = bounds_nr
local_id = gi_global - (worldid * ngeom)
# TODO: Investigate branch elimination with static loop unrolling
if gtype == GeomType.PLANE:
d, n = ray_plane(
geom_xpos_in[worldid, gi],
geom_xmat_in[worldid, gi],
geom_size[worldid % geom_size.shape[0], gi],
ray_origin_world,
ray_dir_world,
)
if gtype == GeomType.HFIELD:
d, n, u, v, f, geom_hfield_id = ray_mesh_with_bvh(
hfield_bvh_id,
geom_dataid[worldid % geom_dataid.shape[0], gi],
geom_xpos_in[worldid, gi],
geom_xmat_in[worldid, gi],
ray_origin_world,
ray_dir_world,
dist,
cull_backfaces,
)
if gtype == GeomType.SPHERE:
d, n = ray_sphere(
geom_xpos_in[worldid, gi],
geom_size[worldid % geom_size.shape[0], gi][0] * geom_size[worldid % geom_size.shape[0], gi][0],
ray_origin_world,
ray_dir_world,
)
if gtype == GeomType.ELLIPSOID:
d, n = ray_ellipsoid(
geom_xpos_in[worldid, gi],
geom_xmat_in[worldid, gi],
geom_size[worldid % geom_size.shape[0], gi],
ray_origin_world,
ray_dir_world,
)
if gtype == GeomType.CAPSULE:
d, n = ray_capsule(
geom_xpos_in[worldid, gi],
geom_xmat_in[worldid, gi],
geom_size[worldid % geom_size.shape[0], gi],
ray_origin_world,
ray_dir_world,
)
if gtype == GeomType.CYLINDER:
d, n = ray_cylinder(
geom_xpos_in[worldid, gi],
geom_xmat_in[worldid, gi],
geom_size[worldid % geom_size.shape[0], gi],
ray_origin_world,
ray_dir_world,
)
if gtype == GeomType.BOX:
d, all, n = ray_box(
geom_xpos_in[worldid, gi],
geom_xmat_in[worldid, gi],
geom_size[worldid % geom_size.shape[0], gi],
ray_origin_world,
ray_dir_world,
)
if gtype == GeomType.MESH:
d, n, u, v, f, hit_mesh_id = ray_mesh_with_bvh(
mesh_bvh_id,
geom_dataid[worldid % geom_dataid.shape[0], gi],
geom_xpos_in[worldid, gi],
geom_xmat_in[worldid, gi],
ray_origin_world,
ray_dir_world,
dist,
cull_backfaces,
)
if gtype == GeomType.FLEX:
hit_geom_id = -2
flexid = flex_geom_flexid[gi]
edge_id = flex_geom_edgeid[gi]
d = float(-1.0)
hit_mesh_id = int(-1)
u = float(0.0)
v = float(0.0)
f = int(-1)
n = wp.vec3(0.0, 0.0, 0.0)
hit_geom_id = int(-1)
if edge_id >= 0:
edge = flex_edge[edge_id]
vert_adr = flex_vertadr[flexid]
v0 = flexvert_xpos_in[worldid, vert_adr + edge[0]]
v1 = flexvert_xpos_in[worldid, vert_adr + edge[1]]
pos = 0.5 * (v0 + v1)
vec = v1 - v0
length = wp.length(vec)
edgeq = math.quat_z2vec(vec)
mat = math.quat_to_mat(edgeq)
size = wp.vec3(flex_radius[flexid], 0.5 * length, 0.0)
d, n = ray_capsule(pos, mat, size, ray_origin_world, ray_dir_world)
hit_mesh_id = flexid
if local_id < bvh_ngeom:
gi = enabled_geom_ids[local_id]
gtype = geom_type[gi]
else:
flex_gr = flex_group_root[worldid, flexid]
d, n, u, v, f = ray_flex_with_bvh(flex_bvh_id, flexid, flex_gr, ray_origin_world, ray_dir_world, dist)
if d >= 0.0:
hit_mesh_id = flexid
gi = local_id - bvh_ngeom
gtype = GeomType.FLEX
# Backface cull: drop exit-face hits when the ray origin is inside the geom,
# matching ray_mesh_with_bvh's `dot(lvec, n) < 0` rule.
if cull_backfaces and d >= 0.0 and wp.dot(ray_dir_world, n) > 0.0:
d = -1.0
hit_geom_id = gi
if d >= 0.0 and d < dist:
dist = d
normal = n
geom_id = hit_geom_id
bary_u = u
bary_v = v
face_idx = f
geom_mesh_id = hit_mesh_id
if wp.static(int(GeomType.PLANE) in geom_ray_types):
if gtype == GeomType.PLANE:
d, n = ray_plane(
geom_xpos_in[worldid, gi],
geom_xmat_in[worldid, gi],
geom_size[worldid % geom_size.shape[0], gi],
ray_origin_world,
ray_dir_world,
)
if wp.static(int(GeomType.HFIELD) in geom_ray_types):
if gtype == GeomType.HFIELD:
d, n, u, v, f, geom_hfield_id = ray_mesh_with_bvh(
hfield_bvh_id,
geom_dataid[worldid % geom_dataid.shape[0], gi],
geom_xpos_in[worldid, gi],
geom_xmat_in[worldid, gi],
ray_origin_world,
ray_dir_world,
dist,
cull_backfaces,
)
if wp.static(int(GeomType.SPHERE) in geom_ray_types):
if gtype == GeomType.SPHERE:
d, n = ray_sphere(
geom_xpos_in[worldid, gi],
geom_size[worldid % geom_size.shape[0], gi][0] * geom_size[worldid % geom_size.shape[0], gi][0],
ray_origin_world,
ray_dir_world,
)
if wp.static(int(GeomType.ELLIPSOID) in geom_ray_types):
if gtype == GeomType.ELLIPSOID:
d, n = ray_ellipsoid(
geom_xpos_in[worldid, gi],
geom_xmat_in[worldid, gi],
geom_size[worldid % geom_size.shape[0], gi],
ray_origin_world,
ray_dir_world,
)
if wp.static(int(GeomType.CAPSULE) in geom_ray_types):
if gtype == GeomType.CAPSULE:
d, n = ray_capsule(
geom_xpos_in[worldid, gi],
geom_xmat_in[worldid, gi],
geom_size[worldid % geom_size.shape[0], gi],
ray_origin_world,
ray_dir_world,
)
if wp.static(int(GeomType.CYLINDER) in geom_ray_types):
if gtype == GeomType.CYLINDER:
d, n = ray_cylinder(
geom_xpos_in[worldid, gi],
geom_xmat_in[worldid, gi],
geom_size[worldid % geom_size.shape[0], gi],
ray_origin_world,
ray_dir_world,
)
if wp.static(int(GeomType.BOX) in geom_ray_types):
if gtype == GeomType.BOX:
d, all, n = ray_box(
geom_xpos_in[worldid, gi],
geom_xmat_in[worldid, gi],
geom_size[worldid % geom_size.shape[0], gi],
ray_origin_world,
ray_dir_world,
)
if wp.static(int(GeomType.MESH) in geom_ray_types):
if gtype == GeomType.MESH:
if wp.static(first_hit):
hit = ray_mesh_with_bvh_anyhit(
mesh_bvh_id,
geom_dataid[worldid % geom_dataid.shape[0], gi],
geom_xpos_in[worldid, gi],
geom_xmat_in[worldid, gi],
ray_origin_world,
ray_dir_world,
dist,
)
d = 0.0 if hit else -1.0
else:
d, n, u, v, f, hit_mesh_id = ray_mesh_with_bvh(
mesh_bvh_id,
geom_dataid[worldid % geom_dataid.shape[0], gi],
geom_xpos_in[worldid, gi],
geom_xmat_in[worldid, gi],
ray_origin_world,
ray_dir_world,
dist,
cull_backfaces,
)
if wp.static(int(GeomType.FLEX) in geom_ray_types):
if gtype == GeomType.FLEX:
hit_geom_id = -2
flexid = flex_geom_flexid[gi]
edge_id = flex_geom_edgeid[gi]
return geom_id, dist, normal, bary_u, bary_v, face_idx, geom_mesh_id
if edge_id >= 0:
edge = flex_edge[edge_id]
vert_adr = flex_vertadr[flexid]
v0 = flexvert_xpos_in[worldid, vert_adr + edge[0]]
v1 = flexvert_xpos_in[worldid, vert_adr + edge[1]]
pos = 0.5 * (v0 + v1)
vec = v1 - v0
length = wp.length(vec)
edgeq = math.quat_z2vec(vec)
mat = math.quat_to_mat(edgeq)
size = wp.vec3(flex_radius[flexid], 0.5 * length, 0.0)
@wp.func
def cast_ray_first_hit(
# Model:
geom_type: wp.array[int],
geom_dataid: wp.array2d[int],
geom_size: wp.array2d[wp.vec3],
flex_vertadr: wp.array[int],
flex_edge: wp.array[wp.vec2i],
flex_radius: wp.array[float],
# Data in:
geom_xpos_in: wp.array2d[wp.vec3],
geom_xmat_in: wp.array2d[wp.mat33],
flexvert_xpos_in: wp.array2d[wp.vec3],
# In:
bvh_id: wp.uint64,
group_root: int,
worldid: int,
bvh_ngeom: int,
bvh_nflexgeom: int,
enabled_geom_ids: wp.array[int],
mesh_bvh_id: wp.array[wp.uint64],
hfield_bvh_id: wp.array[wp.uint64],
flex_geom_flexid: wp.array[int],
flex_geom_edgeid: wp.array[int],
flex_bvh_id: wp.array[wp.uint64],
flex_group_root: wp.array2d[int],
ray_origin_world: wp.vec3,
ray_dir_world: wp.vec3,
max_dist: float,
cull_backfaces: bool,
) -> bool:
"""A simpler version of casting rays that only checks for the first hit."""
query = wp.bvh_query_ray(bvh_id, ray_origin_world, ray_dir_world, group_root)
bounds_nr = int(0)
ngeom = bvh_ngeom + bvh_nflexgeom
d, n = ray_capsule(pos, mat, size, ray_origin_world, ray_dir_world)
hit_mesh_id = flexid
else:
if wp.static(first_hit):
hit = ray_flex_with_bvh_anyhit(
flex_bvh_id,
flexid,
flex_group_root[worldid, flexid],
ray_origin_world,
ray_dir_world,
dist,
)
d = 0.0 if hit else -1.0
else:
flex_gr = flex_group_root[worldid, flexid]
d, n, u, v, f = ray_flex_with_bvh(flex_bvh_id, flexid, flex_gr, ray_origin_world, ray_dir_world, dist)
if d >= 0.0:
hit_mesh_id = flexid
while wp.bvh_query_next(query, bounds_nr, max_dist):
gi_global = bounds_nr
local_id = gi_global - (worldid * ngeom)
# Backface cull: drop exit-face hits when the ray origin is inside the geom,
# matching ray_mesh_with_bvh's `dot(lvec, n) < 0` rule. Strict `> 0` keeps
# tangent hits and skips branches with a zero-vector normal (any-hit).
if cull_backfaces and d >= 0.0 and wp.dot(ray_dir_world, n) > 0.0:
d = -1.0
d = float(-1.0)
n = wp.vec3(0.0, 0.0, 0.0)
if local_id < bvh_ngeom:
gi = enabled_geom_ids[local_id]
gtype = geom_type[gi]
else:
gi = local_id - bvh_ngeom
gtype = GeomType.FLEX
# TODO: Investigate branch elimination with static loop unrolling
if gtype == GeomType.PLANE:
d, n = ray_plane(
geom_xpos_in[worldid, gi],
geom_xmat_in[worldid, gi],
geom_size[worldid % geom_size.shape[0], gi],
ray_origin_world,
ray_dir_world,
)
if gtype == GeomType.HFIELD:
d, n, u, v, f, geom_hfield_id = ray_mesh_with_bvh(
hfield_bvh_id,
geom_dataid[worldid % geom_dataid.shape[0], gi],
geom_xpos_in[worldid, gi],
geom_xmat_in[worldid, gi],
ray_origin_world,
ray_dir_world,
max_dist,
cull_backfaces,
)
if gtype == GeomType.SPHERE:
d, n = ray_sphere(
geom_xpos_in[worldid, gi],
geom_size[worldid % geom_size.shape[0], gi][0] * geom_size[worldid % geom_size.shape[0], gi][0],
ray_origin_world,
ray_dir_world,
)
if gtype == GeomType.ELLIPSOID:
d, n = ray_ellipsoid(
geom_xpos_in[worldid, gi],
geom_xmat_in[worldid, gi],
geom_size[worldid % geom_size.shape[0], gi],
ray_origin_world,
ray_dir_world,
)
if gtype == GeomType.CAPSULE:
d, n = ray_capsule(
geom_xpos_in[worldid, gi],
geom_xmat_in[worldid, gi],
geom_size[worldid % geom_size.shape[0], gi],
ray_origin_world,
ray_dir_world,
)
if gtype == GeomType.CYLINDER:
d, n = ray_cylinder(
geom_xpos_in[worldid, gi],
geom_xmat_in[worldid, gi],
geom_size[worldid % geom_size.shape[0], gi],
ray_origin_world,
ray_dir_world,
)
if gtype == GeomType.BOX:
d, all, n = ray_box(
geom_xpos_in[worldid, gi],
geom_xmat_in[worldid, gi],
geom_size[worldid % geom_size.shape[0], gi],
ray_origin_world,
ray_dir_world,
)
if gtype == GeomType.MESH:
hit = ray_mesh_with_bvh_anyhit(
mesh_bvh_id,
geom_dataid[worldid % geom_dataid.shape[0], gi],
geom_xpos_in[worldid, gi],
geom_xmat_in[worldid, gi],
ray_origin_world,
ray_dir_world,
max_dist,
)
d = 0.0 if hit else -1.0
if gtype == GeomType.FLEX:
flexid = flex_geom_flexid[gi]
edge_id = flex_geom_edgeid[gi]
if edge_id >= 0:
edge = flex_edge[edge_id]
vert_adr = flex_vertadr[flexid]
v0 = flexvert_xpos_in[worldid, vert_adr + edge[0]]
v1 = flexvert_xpos_in[worldid, vert_adr + edge[1]]
pos = 0.5 * (v0 + v1)
vec = v1 - v0
length = wp.length(vec)
edgeq = math.quat_z2vec(vec)
mat = math.quat_to_mat(edgeq)
size = wp.vec3(flex_radius[flexid], 0.5 * length, 0.0)
d, n = ray_capsule(pos, mat, size, ray_origin_world, ray_dir_world)
if wp.static(first_hit):
# Any-hit: return as soon as anything is in range; surface data is unused.
if d >= 0.0 and d < dist:
return hit_geom_id, d, n, u, v, f, hit_mesh_id
else:
hit = ray_flex_with_bvh_anyhit(
flex_bvh_id,
flexid,
flex_group_root[worldid, flexid],
ray_origin_world,
ray_dir_world,
max_dist,
)
d = 0.0 if hit else -1.0
if d >= 0.0 and d < dist:
dist = d
normal = n
geom_id = hit_geom_id
bary_u = u
bary_v = v
face_idx = f
geom_mesh_id = hit_mesh_id
# Backface cull: see cast_ray for rationale. Strict `> 0` keeps tangent
# hits and skips branches with a zero-vector normal (mesh/flex anyhit).
if cull_backfaces and d >= 0.0 and wp.dot(ray_dir_world, n) > 0.0:
d = -1.0
return geom_id, dist, normal, bary_u, bary_v, face_idx, geom_mesh_id
if d >= 0.0 and d < max_dist:
return True
return False
return cast_ray
@wp.func
def compute_lighting(
# Model:
geom_type: wp.array[int],
geom_dataid: wp.array2d[int],
geom_size: wp.array2d[wp.vec3],
flex_vertadr: wp.array[int],
flex_edge: wp.array[wp.vec2i],
flex_radius: wp.array[float],
# Data in:
geom_xpos_in: wp.array2d[wp.vec3],
geom_xmat_in: wp.array2d[wp.mat33],
flexvert_xpos_in: wp.array2d[wp.vec3],
# In:
use_shadows: bool,
bvh_id: wp.uint64,
group_root: int,
bvh_ngeom: int,
bvh_nflexgeom: int,
enabled_geom_ids: wp.array[int],
worldid: int,
mesh_bvh_id: wp.array[wp.uint64],
hfield_bvh_id: wp.array[wp.uint64],
flex_geom_flexid: wp.array[int],
flex_geom_edgeid: wp.array[int],
flex_bvh_id: wp.array[wp.uint64],
flex_group_root: wp.array2d[int],
lightactive: bool,
lighttype: int,
lightcastshadow: bool,
lightpos: wp.vec3,
lightdir: wp.vec3,
normal: wp.vec3,
hitpoint: wp.vec3,
cull_backfaces: bool,
) -> float:
light_contribution = float(0.0)
def _make_compute_lighting(cast_ray_first_hit: wp.Function) -> wp.Function:
"""Build specialized compute_lighting."""
# TODO: We should probably only be looping over active lights
# in the first place with a static loop of enabled light idx?
if not lightactive:
return light_contribution
@wp.func
def compute_lighting(
# Model:
geom_type: wp.array[int],
geom_dataid: wp.array2d[int],
geom_size: wp.array2d[wp.vec3],
flex_vertadr: wp.array[int],
flex_edge: wp.array[wp.vec2i],
flex_radius: wp.array[float],
# Data in:
geom_xpos_in: wp.array2d[wp.vec3],
geom_xmat_in: wp.array2d[wp.mat33],
flexvert_xpos_in: wp.array2d[wp.vec3],
# In:
use_shadows: bool,
bvh_id: wp.uint64,
group_root: int,
bvh_ngeom: int,
bvh_nflexgeom: int,
enabled_geom_ids: wp.array[int],
worldid: int,
mesh_bvh_id: wp.array[wp.uint64],
hfield_bvh_id: wp.array[wp.uint64],
flex_geom_flexid: wp.array[int],
flex_geom_edgeid: wp.array[int],
flex_bvh_id: wp.array[wp.uint64],
flex_group_root: wp.array2d[int],
lightactive: bool,
lighttype: int,
lightcastshadow: bool,
lightpos: wp.vec3,
lightdir: wp.vec3,
normal: wp.vec3,
hitpoint: wp.vec3,
cull_backfaces: bool,
) -> float:
light_contribution = float(0.0)
L = wp.vec3(0.0, 0.0, 0.0)
dist_to_light = float(MJ_MAXVAL)
attenuation = float(1.0)
# TODO: We should probably only be looping over active lights
# in the first place with a static loop of enabled light idx?
if not lightactive:
return light_contribution
if lighttype == 1: # directional light
L = wp.normalize(-lightdir)
else:
L, dist_to_light = math.normalize_with_norm(lightpos - hitpoint)
attenuation = 1.0 / (1.0 + 0.02 * dist_to_light * dist_to_light)
if lighttype == 0: # spot light
spot_dir = wp.normalize(lightdir)
cos_theta = wp.dot(-L, spot_dir)
spot_factor = wp.min(1.0, wp.max(0.0, (cos_theta - 0.85) / (0.95 - 0.85)))
attenuation = attenuation * spot_factor
L = wp.vec3(0.0, 0.0, 0.0)
dist_to_light = float(MJ_MAXVAL)
attenuation = float(1.0)
ndotl = wp.max(0.0, wp.dot(normal, L))
if ndotl == 0.0:
return light_contribution
visible = float(1.0)
if use_shadows and lightcastshadow:
# Nudge the origin slightly along the surface normal to avoid
# self-intersection when casting shadow rays
eps = 1.0e-4
shadow_origin = hitpoint + normal * eps
# Distance-limited shadows: cap by dist_to_light (for non-directional)
max_t = float(dist_to_light - 1.0e-3)
if lighttype == 1: # directional light
max_t = float(1.0e8)
L = wp.normalize(-lightdir)
else:
L, dist_to_light = math.normalize_with_norm(lightpos - hitpoint)
attenuation = 1.0 / (1.0 + 0.02 * dist_to_light * dist_to_light)
if lighttype == 0: # spot light
spot_dir = wp.normalize(lightdir)
cos_theta = wp.dot(-L, spot_dir)
spot_factor = wp.min(1.0, wp.max(0.0, (cos_theta - 0.85) * 10.0))
attenuation = attenuation * spot_factor
shadow_hit = cast_ray_first_hit(
geom_type,
geom_dataid,
geom_size,
flex_vertadr,
flex_edge,
flex_radius,
geom_xpos_in,
geom_xmat_in,
flexvert_xpos_in,
bvh_id,
group_root,
worldid,
bvh_ngeom,
bvh_nflexgeom,
enabled_geom_ids,
mesh_bvh_id,
hfield_bvh_id,
flex_geom_flexid,
flex_geom_edgeid,
flex_bvh_id,
flex_group_root,
shadow_origin,
L,
max_t,
cull_backfaces,
)
ndotl = wp.max(0.0, wp.dot(normal, L))
if ndotl == 0.0:
return light_contribution
if shadow_hit:
visible = 0.3
visible = float(1.0)
return ndotl * attenuation * visible
if use_shadows and lightcastshadow:
# Nudge the origin slightly along the surface normal to avoid
# self-intersection when casting shadow rays
shadow_origin = hitpoint + normal * 1.0e-4
# Distance-limited shadows: cap by dist_to_light (for non-directional)
max_t = dist_to_light - 1.0e-3
if lighttype == 1: # directional light
max_t = 1.0e8
shadow_geom_id, shadow_d, shadow_n, shadow_u, shadow_v, shadow_f, shadow_mesh_id = cast_ray_first_hit(
geom_type,
geom_dataid,
geom_size,
flex_vertadr,
flex_edge,
flex_radius,
geom_xpos_in,
geom_xmat_in,
flexvert_xpos_in,
bvh_id,
group_root,
worldid,
bvh_ngeom,
bvh_nflexgeom,
enabled_geom_ids,
mesh_bvh_id,
hfield_bvh_id,
flex_geom_flexid,
flex_geom_edgeid,
flex_bvh_id,
flex_group_root,
shadow_origin,
L,
max_t,
cull_backfaces,
)
if shadow_geom_id != -1:
visible = 0.3
return ndotl * attenuation * visible
return compute_lighting
@event_scope
@@ -611,6 +508,13 @@ def render(m: Model, d: Data, rc: RenderContext):
rc.depth_data.fill_(0.0)
rc.seg_data.fill_(wp.vec2i(-1, -1))
# Specialize the ray-cast helpers to the geom types present in the scene so the
# compiler eliminates intersection branches for absent types.
geom_ray_types = rc.geom_ray_types
cast_ray = _make_cast_ray(geom_ray_types, first_hit=False)
cast_ray_first_hit = _make_cast_ray(geom_ray_types, first_hit=True)
compute_lighting = _make_compute_lighting(cast_ray_first_hit)
@wp.kernel(module="unique", enable_backward=False)
def _render_megakernel(
# Model:
@@ -676,31 +580,31 @@ def render(m: Model, d: Data, rc: RenderContext):
):
worldid, rayid = wp.tid()
# Map global rayid -> (cam_idx, rayid_local) using cumulative sizes
cam_idx = int(-1)
# Map global rayid -> (camid, rayid_local) using cumulative sizes
camid = int(-1)
rayid_local = int(-1)
accum = int(0)
for i in range(nrender):
num_i = cam_res[i][0] * cam_res[i][1]
if rayid < accum + num_i:
cam_idx = i
camid = i
rayid_local = rayid - accum
break
accum += num_i
if cam_idx == -1 or rayid_local < 0:
if camid == -1 or rayid_local < 0:
return
if not render_rgb[cam_idx] and not render_depth[cam_idx] and not render_seg[cam_idx]:
if not render_rgb[camid] and not render_depth[camid] and not render_seg[camid]:
return
# Map active camera index to MuJoCo camera ID
mujoco_cam_id = cam_id_map[cam_idx]
mujoco_cam_id = cam_id_map[camid]
if wp.static(rc.use_precomputed_rays):
ray_dir_local_cam = ray[rayid]
else:
img_w = cam_res[cam_idx][0]
img_h = cam_res[cam_idx][1]
img_w = cam_res[camid][0]
img_h = cam_res[camid][1]
px = rayid_local % img_w
py = rayid_local // img_w
ray_dir_local_cam = compute_ray(
@@ -742,24 +646,25 @@ def render(m: Model, d: Data, rc: RenderContext):
flex_group_root,
ray_origin_world,
ray_dir_world,
float(MJ_MAXVAL),
wp.static(rc.enable_backface_culling),
)
if render_seg[cam_idx] and geom_id != -1:
if render_seg[camid] and geom_id != -1:
if geom_id == -2:
seg_out[worldid, seg_adr[cam_idx] + rayid_local] = wp.vec2i(mesh_id, int(ObjType.FLEX))
seg_out[worldid, seg_adr[camid] + rayid_local] = wp.vec2i(mesh_id, int(ObjType.FLEX))
else:
seg_out[worldid, seg_adr[cam_idx] + rayid_local] = wp.vec2i(geom_id, int(ObjType.GEOM))
seg_out[worldid, seg_adr[camid] + rayid_local] = wp.vec2i(geom_id, int(ObjType.GEOM))
# Early Out
if geom_id == -1:
if wp.static(rc.render_skybox) and render_rgb[cam_idx]:
if wp.static(rc.render_skybox) and render_rgb[camid]:
skybox_color = sample_skybox(
textures[wp.static(rc.skybox_tex_id)],
wp.static(1.0 / float(rc.skybox_face_width)),
ray_dir_world,
)
rgb_out[worldid, rgb_adr[cam_idx] + rayid_local] = pack_rgba_to_uint32(
rgb_out[worldid, rgb_adr[camid] + rayid_local] = pack_rgba_to_uint32(
skybox_color[0] * 255.0,
skybox_color[1] * 255.0,
skybox_color[2] * 255.0,
@@ -767,14 +672,14 @@ def render(m: Model, d: Data, rc: RenderContext):
)
return
if render_depth[cam_idx]:
if render_depth[camid]:
# Planar depth: project Euclidean distance onto the camera's optical axis.
# In camera-local coordinates, the optical axis is -Z. The Z-component of the
# normalized ray direction is negative, so -ray_dir_local_cam[2] gives cos(θ)
# between the ray and the optical axis.
depth_out[worldid, depth_adr[cam_idx] + rayid_local] = dist * (-ray_dir_local_cam[2])
depth_out[worldid, depth_adr[camid] + rayid_local] = dist * (-ray_dir_local_cam[2])
if not render_rgb[cam_idx]:
if not render_rgb[camid]:
return
# Shade the pixel
@@ -864,7 +769,7 @@ def render(m: Model, d: Data, rc: RenderContext):
hit_color = wp.min(result, wp.vec3(1.0, 1.0, 1.0))
hit_color = wp.max(hit_color, wp.vec3(0.0, 0.0, 0.0))
rgb_out[worldid, rgb_adr[cam_idx] + rayid_local] = pack_rgba_to_uint32(
rgb_out[worldid, rgb_adr[camid] + rayid_local] = pack_rgba_to_uint32(
hit_color[0] * 255.0,
hit_color[1] * 255.0,
hit_color[2] * 255.0,
@@ -934,4 +839,5 @@ def render(m: Model, d: Data, rc: RenderContext):
rc.depth_data,
rc.seg_data,
],
block_dim=m.block_dim.render,
)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py
Vendored
+3 -3
View File
@@ -1012,7 +1012,7 @@ def get_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt
elif element == State.EQ_ACTIVE:
for j in range(neq):
state_out[worldid, adr + j] = float(eq_active_in[worldid, j])
adr += j
adr += neq
elif element == State.MOCAP_POS:
for j in range(nmocap):
pos = mocap_pos_in[worldid, j]
@@ -1160,12 +1160,12 @@ def set_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt
elif element == State.EQ_ACTIVE:
for j in range(neq):
eq_active_out[worldid, j] = bool(state_in[worldid, adr + j])
adr += j
adr += neq
elif element == State.MOCAP_POS:
for j in range(nmocap):
pos = wp.vec3(
state_in[worldid, adr + 1],
state_in[worldid, adr + 0],
state_in[worldid, adr + 1],
state_in[worldid, adr + 2],
)
mocap_pos_out[worldid, j] = pos
mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py
Vendored
+6
View File
@@ -67,6 +67,7 @@ class BlockDim:
linesearch_iterative: linesearch iterative block dimension (solver)
contact_jac_tiled: contact Jacobian tiled block dimension (solver)
qderiv_actuator_dense: qderiv actuator dense block dimension (derivative)
render: render block dimension (render)
"""
# collision_driver
@@ -93,6 +94,8 @@ class BlockDim:
contact_jac_tiled: int = 32
# derivative
qderiv_actuator_dense: int = 32
# render
render: int = 64
class BroadphaseType(enum.IntEnum):
@@ -2206,6 +2209,8 @@ class RenderContext:
mesh-ray rule. When False, the renderer reports inner-surface hits, which
is faster but causes a camera placed inside a geom to render that geom's
back wall.
geom_ray_types: tuple of GeomType int values present in the scene, used to
statically eliminate unused intersection branches in the ray-cast kernels.
"""
nrender: int
@@ -2260,3 +2265,4 @@ class RenderContext:
znear: float
total_rays: int
enable_backface_culling: bool
geom_ray_types: tuple = ()
mjx/mujoco/mjx/warp/bvh.py
+4 -7
View File
@@ -14,19 +14,17 @@
# ==============================================================================
"""DO NOT EDIT. This file is auto-generated."""
import dataclasses
import functools
import jax
import warp as wp
from mujoco.mjx._src import types
import mujoco.mjx.third_party.mujoco_warp as mjwarp
from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
from mujoco.mjx.warp import ffi
from mujoco.mjx.warp.render_context import _MJX_RENDER_CONTEXT_BUFFERS
from mujoco.mjx.warp.render_context import RenderContextPytree
import mujoco.mjx.third_party.mujoco_warp as mjwarp
from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
import warp as wp
_m = mjwarp.Model(
**{f.name: None for f in dataclasses.fields(mjwarp.Model) if f.init}
@@ -50,7 +48,6 @@ _cb = mjwp_types.Callback(
**{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init}
)
@ffi.format_args_for_warp
def _refit_bvh_shim(
# Model
mjx/mujoco/mjx/warp/collision_driver.py
+2 -6
View File
@@ -14,17 +14,14 @@
# ==============================================================================
"""DO NOT EDIT. This file is auto-generated."""
import dataclasses
import functools
import jax
import warp as wp
from mujoco.mjx._src import types
from mujoco.mjx.warp import ffi
import mujoco.mjx.third_party.mujoco_warp as mjwarp
from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
from mujoco.mjx.warp import ffi
import warp as wp
_m = mjwarp.Model(
**{f.name: None for f in dataclasses.fields(mjwarp.Model) if f.init}
@@ -48,7 +45,6 @@ _cb = mjwp_types.Callback(
**{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init}
)
@ffi.format_args_for_warp
def _collision_shim(
# Model
mjx/mujoco/mjx/warp/forward.py
+2 -6
View File
@@ -14,17 +14,14 @@
# ==============================================================================
"""DO NOT EDIT. This file is auto-generated."""
import dataclasses
import functools
import jax
import warp as wp
from mujoco.mjx._src import types
from mujoco.mjx.warp import ffi
import mujoco.mjx.third_party.mujoco_warp as mjwarp
from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
from mujoco.mjx.warp import ffi
import warp as wp
_m = mjwarp.Model(
**{f.name: None for f in dataclasses.fields(mjwarp.Model) if f.init}
@@ -48,7 +45,6 @@ _cb = mjwp_types.Callback(
**{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init}
)
@ffi.format_args_for_warp
def _forward_shim(
# Model
mjx/mujoco/mjx/warp/render.py
+7 -6
View File
@@ -14,19 +14,17 @@
# ==============================================================================
"""DO NOT EDIT. This file is auto-generated."""
import dataclasses
import functools
import jax
import warp as wp
from mujoco.mjx._src import types
import mujoco.mjx.third_party.mujoco_warp as mjwarp
from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
from mujoco.mjx.warp import ffi
from mujoco.mjx.warp.render_context import _MJX_RENDER_CONTEXT_BUFFERS
from mujoco.mjx.warp.render_context import RenderContextPytree
import mujoco.mjx.third_party.mujoco_warp as mjwarp
from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
import warp as wp
_m = mjwarp.Model(
**{f.name: None for f in dataclasses.fields(mjwarp.Model) if f.init}
@@ -55,6 +53,7 @@ _cb = mjwp_types.Callback(
def _render_shim(
# Model
nworld: int,
block_dim: mjwp_types.BlockDim,
cam_fovy: wp.array2d[float],
cam_intrinsic: wp.array2d[wp.vec4],
cam_projection: wp.array[int],
@@ -94,6 +93,7 @@ def _render_shim(
_m.callback = _cb
_d.efc = _e
_d.contact = _c
_m.block_dim = block_dim
_m.cam_fovy = cam_fovy
_m.cam_intrinsic = cam_intrinsic
_m.cam_projection = cam_projection
@@ -164,6 +164,7 @@ def _render_jax_impl(m: types.Model, d: types.Data, ctx: RenderContextPytree):
)
out = jf(
render_ctx.nworld,
m._impl.block_dim,
m.cam_fovy,
m.cam_intrinsic,
m._impl.cam_projection,
mjx/mujoco/mjx/warp/smooth.py
+2 -6
View File
@@ -14,17 +14,14 @@
# ==============================================================================
"""DO NOT EDIT. This file is auto-generated."""
import dataclasses
import functools
import jax
import warp as wp
from mujoco.mjx._src import types
from mujoco.mjx.warp import ffi
import mujoco.mjx.third_party.mujoco_warp as mjwarp
from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
from mujoco.mjx.warp import ffi
import warp as wp
_m = mjwarp.Model(
**{f.name: None for f in dataclasses.fields(mjwarp.Model) if f.init}
@@ -48,7 +45,6 @@ _cb = mjwp_types.Callback(
**{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init}
)
@ffi.format_args_for_warp
def _kinematics_shim(
# Model
mjx/mujoco/mjx/warp/types.py
+9 -24
View File
@@ -15,17 +15,14 @@
"""MJX Warp types.
DO NOT EDIT. This file is auto-generated.
"""
import dataclasses
import typing
from typing import Tuple
import jax
from jax import tree_util
from jax.interpreters import batching
import numpy as np
from mujoco.mjx._src import dataclasses as mjx_dataclasses
import numpy as np
if typing.TYPE_CHECKING:
GraphMode = int
@@ -37,7 +34,6 @@ if typing.TYPE_CHECKING:
else:
try:
from warp._src.jax_experimental.ffi import GraphMode
from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
Callback = mjwp_types.Callback
@@ -46,13 +42,6 @@ else:
Callback = None
PyTreeNode = mjx_dataclasses.PyTreeNode
def _as_numpy_array(value):
if hasattr(value, 'numpy'):
return value.numpy()
return np.asarray(value)
@dataclasses.dataclass(frozen=True)
@tree_util.register_pytree_node_class
class TileSet:
@@ -64,19 +53,20 @@ class TileSet:
adr: address of each tile in the set
size: size of all the tiles in this set
"""
adr: np.ndarray
size: int
# Manually kept in this generated shim until TileSet method generation is
# needed more broadly. Keep this in sync with mujoco_warp._src.types.TileSet.
def __eq__(self, other) -> bool:
if self.__class__ is not other.__class__:
return NotImplemented
return self.size == other.size and np.array_equal(
_as_numpy_array(self.adr), _as_numpy_array(other.adr)
np.asarray(self.adr), np.asarray(other.adr)
)
def __hash__(self) -> int:
adr = _as_numpy_array(self.adr)
adr = np.asarray(self.adr)
return hash((self.size, adr.dtype.str, adr.shape, adr.tobytes()))
def tree_flatten(self):
@@ -118,8 +108,8 @@ class BlockDim:
linesearch_iterative: linesearch iterative block dimension (solver)
contact_jac_tiled: contact Jacobian tiled block dimension (solver)
qderiv_actuator_dense: qderiv actuator dense block dimension (derivative)
render: render block dimension (render)
"""
actuator_velocity: int
cholesky_factorize: int
cholesky_factorize_solve: int
@@ -131,6 +121,7 @@ class BlockDim:
linesearch_iterative: int
qderiv_actuator_dense: int
ray: int
render: int
segmented_sort: int
solve_LD_sparse_fused: int
update_gradient_JTDAJ_dense: int
@@ -150,13 +141,10 @@ class BlockDim:
class StatisticWarp(PyTreeNode):
"""Derived fields from Statistic."""
meaninertia: jax.Array
class OptionWarp(PyTreeNode):
"""Derived fields from Option."""
broadphase: int
broadphase_filter: int
ccd_iterations: int
@@ -171,7 +159,6 @@ class OptionWarp(PyTreeNode):
sdf_initpoints: int
sdf_iterations: int
class ModelWarp(PyTreeNode):
"""Derived fields from Model."""
D_colind: np.ndarray
@@ -353,7 +340,6 @@ class ModelWarp(PyTreeNode):
wrap_site_adr: np.ndarray
wrap_site_pair_adr: np.ndarray
class DataWarp(PyTreeNode):
"""Derived fields from Data."""
M: jax.Array
@@ -465,8 +451,6 @@ class DataWarp(PyTreeNode):
wrap_obj: jax.Array
wrap_xpos: jax.Array
shape = property(lambda self: self.cacc.shape)
DATA_NON_VMAP = {
'contact__dim',
'contact__dist',
@@ -495,7 +479,6 @@ DATA_NON_VMAP = {
'nworld',
}
def _to_elt(cont, _, d, axis):
return DataWarp(**{
f.name: (
@@ -729,6 +712,7 @@ _NDIM = {
'block_dim__linesearch_iterative': 0,
'block_dim__qderiv_actuator_dense': 0,
'block_dim__ray': 0,
'block_dim__render': 0,
'block_dim__segmented_sort': 0,
'block_dim__solve_LD_sparse_fused': 0,
'block_dim__update_gradient_JTDAJ_dense': 0,
@@ -1364,6 +1348,7 @@ _BATCH_DIM = {
'block_dim__linesearch_iterative': False,
'block_dim__qderiv_actuator_dense': False,
'block_dim__ray': False,
'block_dim__render': False,
'block_dim__segmented_sort': False,
'block_dim__solve_LD_sparse_fused': False,
'block_dim__update_gradient_JTDAJ_dense': False,
plugin/usd_decoder/usd_decoder.cc
+227 -39
View File
@@ -46,6 +46,8 @@
#include <pxr/base/gf/matrix4f.h>
#include <pxr/base/gf/rotation.h>
#include <pxr/base/gf/vec3d.h>
#include <pxr/base/tf/staticData.h>
#include <pxr/base/tf/staticTokens.h>
#include <pxr/base/tf/token.h>
#include <pxr/base/vt/types.h>
#include <pxr/usd/sdf/path.h>
@@ -82,6 +84,27 @@
using pxr::MjcPhysicsTokens;
using pxr::TfToken;
template <typename T>
using TfStaticData = pxr::TfStaticData<T>;
// clang-format off
TF_DEFINE_PRIVATE_TOKENS(kNewtonTokens,
((NewtonMaterialAPI, "NewtonMaterialAPI"))
((NewtonMeshCollisionAPI, "NewtonMeshCollisionAPI"))
((newtonMaxSolverIterations, "newton:maxSolverIterations"))
((newtonTimeStepsPerSecond, "newton:timeStepsPerSecond"))
((newtonGravityEnabled, "newton:gravityEnabled"))
((newtonContactMargin, "newton:contactMargin"))
((newtonContactGap, "newton:contactGap"))
((newtonMaxHullVertices, "newton:maxHullVertices"))
((newtonTorsionalFriction, "newton:torsionalFriction"))
((newtonRollingFriction, "newton:rollingFriction"))
((newtonMimicJoint, "newton:mimicJoint"))
((newtonMimicCoef0, "newton:mimicCoef0"))
((newtonMimicCoef1, "newton:mimicCoef1"))
((NewtonMimicAPI, "NewtonMimicAPI"))
);
// clang-format on
struct UsdCaches {
pxr::UsdGeomXformCache xform_cache;
@@ -463,15 +486,48 @@ void ParseUsdPhysicsScene(mjSpec* spec,
SetGravityAttributes(spec, stage, gravity_direction, gravity_magnitude);
// Early exit if theres no MjcPhysicsSceneAPI applied.
if (!physics_scene.GetPrim().HasAPI<pxr::MjcPhysicsSceneAPI>()) {
// Parse Newton scene attributes if present (works for Newton-only files)
pxr::UsdPrim scene_prim = physics_scene.GetPrim();
auto newton_iterations = scene_prim.GetAttribute(
kNewtonTokens->newtonMaxSolverIterations);
if (newton_iterations && newton_iterations.HasAuthoredValue()) {
int val;
newton_iterations.Get(&val);
if (val >= 0) spec->option.iterations = val;
}
auto newton_timesteps = scene_prim.GetAttribute(
kNewtonTokens->newtonTimeStepsPerSecond);
if (newton_timesteps && newton_timesteps.HasAuthoredValue()) {
int val;
newton_timesteps.Get(&val);
if (val > 0) spec->option.timestep = 1.0 / val;
}
auto newton_gravity = scene_prim.GetAttribute(
kNewtonTokens->newtonGravityEnabled);
if (newton_gravity && newton_gravity.HasAuthoredValue()) {
bool enabled;
newton_gravity.Get(&enabled);
if (!enabled) {
spec->option.disableflags |= mjDSBL_GRAVITY;
}
}
// Early exit if there's no MjcPhysicsSceneAPI applied.
if (!scene_prim.HasAPI<pxr::MjcPhysicsSceneAPI>()) {
return;
}
auto mjc_physics_scene = pxr::MjcPhysicsSceneAPI(physics_scene.GetPrim());
auto mjc_physics_scene = pxr::MjcPhysicsSceneAPI(scene_prim);
double timestep;
mjc_physics_scene.GetTimestepAttr().Get(&timestep);
spec->option.timestep = timestep;
// MJC values override Newton values only when explicitly authored.
auto timestep_attr = mjc_physics_scene.GetTimestepAttr();
if (timestep_attr.HasAuthoredValue()) {
double timestep;
timestep_attr.Get(&timestep);
spec->option.timestep = timestep;
mju_warning("Scene '%s' uses deprecated mjc:option:timestep. "
"Please migrate to newton:timeStepsPerSecond.",
scene_prim.GetPath().GetText());
}
double impratio;
mjc_physics_scene.GetImpRatioAttr().Get(&impratio);
@@ -586,9 +642,15 @@ void ParseUsdPhysicsScene(mjSpec* spec,
spec->option.solver = mjSOL_PGS;
}
int iterations;
mjc_physics_scene.GetIterationsAttr().Get(&iterations);
spec->option.iterations = iterations;
auto iterations_attr = mjc_physics_scene.GetIterationsAttr();
if (iterations_attr.HasAuthoredValue()) {
int iterations;
iterations_attr.Get(&iterations);
spec->option.iterations = iterations;
mju_warning("Scene '%s' uses deprecated mjc:option:iterations. "
"Please migrate to newton:maxSolverIterations.",
scene_prim.GetPath().GetText());
}
int ls_iterations;
mjc_physics_scene.GetLSIterationsAttr().Get(&ls_iterations);
@@ -638,9 +700,19 @@ void ParseUsdPhysicsScene(mjSpec* spec,
mjc_physics_scene.GetDamperFlagAttr().Get(&damper_flag);
spec->option.disableflags |= (!damper_flag ? mjDSBL_DAMPER : 0);
bool gravity_flag;
mjc_physics_scene.GetGravityFlagAttr().Get(&gravity_flag);
spec->option.disableflags |= (!gravity_flag ? mjDSBL_GRAVITY : 0);
auto gravity_flag_attr = mjc_physics_scene.GetGravityFlagAttr();
if (gravity_flag_attr.HasAuthoredValue()) {
bool gravity_flag;
gravity_flag_attr.Get(&gravity_flag);
if (!gravity_flag) {
spec->option.disableflags |= mjDSBL_GRAVITY;
} else {
spec->option.disableflags &= ~mjDSBL_GRAVITY;
}
mju_warning("Scene '%s' uses deprecated mjc:flag:gravity. "
"Please migrate to newton:gravityEnabled.",
scene_prim.GetPath().GetText());
}
bool clampctrl_flag;
mjc_physics_scene.GetClampCtrlFlagAttr().Get(&clampctrl_flag);
@@ -914,13 +986,39 @@ void ParseMjcPhysicsCollisionAPI(
}
auto margin_attr = collision_api.GetMarginAttr();
if (margin_attr.HasAuthoredValue()) {
auto gap_attr = collision_api.GetGapAttr();
bool mjc_margin_authored = margin_attr.HasAuthoredValue();
bool mjc_gap_authored = gap_attr.HasAuthoredValue();
if (mjc_margin_authored) {
margin_attr.Get(&geom->margin);
mju_warning("Prim '%s' uses deprecated mjc:margin. "
"Please migrate to newton:contactMargin and newton:contactGap.",
collision_api.GetPrim().GetPath().GetText());
}
if (mjc_gap_authored) {
gap_attr.Get(&geom->gap);
mju_warning("Prim '%s' uses deprecated mjc:gap. "
"Please migrate to newton:contactGap.",
collision_api.GetPrim().GetPath().GetText());
}
auto gap_attr = collision_api.GetGapAttr();
if (gap_attr.HasAuthoredValue()) {
gap_attr.Get(&geom->gap);
// Newton collision fallback: newton:contactMargin + newton:contactGap -> margin, gap
if (!mjc_margin_authored || !mjc_gap_authored) {
pxr::UsdPrim prim = collision_api.GetPrim();
auto newton_margin = prim.GetAttribute(kNewtonTokens->newtonContactMargin);
auto newton_gap = prim.GetAttribute(kNewtonTokens->newtonContactGap);
float n_margin = 0, n_gap = 0;
bool has_newton_margin = newton_margin && newton_margin.HasAuthoredValue();
bool has_newton_gap = newton_gap && newton_gap.HasAuthoredValue();
if (has_newton_margin) newton_margin.Get(&n_margin);
if (has_newton_gap) newton_gap.Get(&n_gap);
if (!mjc_gap_authored && has_newton_gap) {
geom->gap = n_gap;
}
if (!mjc_margin_authored && has_newton_margin) {
geom->margin = n_margin;
}
}
}
@@ -942,8 +1040,19 @@ void ParseMjcPhysicsMeshCollisionAPI(
}
auto maxhullvert_attr = mesh_collision_api.GetMaxHullVertAttr();
auto newton_maxhull = mesh_collision_api.GetPrim().GetAttribute(
kNewtonTokens->newtonMaxHullVertices);
if (maxhullvert_attr.HasAuthoredValue()) {
maxhullvert_attr.Get(&mesh->maxhullvert);
if (!newton_maxhull || !newton_maxhull.HasAuthoredValue()) {
mju_warning("Prim '%s' uses deprecated mjc:maxhullvert. "
"Please migrate to newton:maxHullVertices.",
mesh_collision_api.GetPrim().GetPath().GetText());
}
} else if (newton_maxhull && newton_maxhull.HasAuthoredValue()) {
int val;
newton_maxhull.Get(&val);
mesh->maxhullvert = val;
}
}
@@ -1668,15 +1777,42 @@ void ParseUsdPhysicsMaterialAPI(
}
void ParseMjcPhysicsMaterialAPI(
mjsGeom* geom, const pxr::MjcPhysicsMaterialAPI& material_api) {
auto torsional_friction_attr = material_api.GetTorsionalFrictionAttr();
if (torsional_friction_attr.HasAuthoredValue()) {
torsional_friction_attr.Get(&geom->friction[1]);
mjsGeom* geom, const pxr::UsdPrim& material_prim,
const pxr::MjcPhysicsMaterialAPI& material_api) {
// Torsional friction: prefer newton:torsionalFriction, fall back to
// mjc:torsionalfriction with deprecation warning. If both are authored,
// mjc takes precedence for backwards compatibility.
auto mjc_torsional = material_api.GetTorsionalFrictionAttr();
auto newton_torsional = material_prim.GetAttribute(
kNewtonTokens->newtonTorsionalFriction);
if (mjc_torsional.HasAuthoredValue()) {
mjc_torsional.Get(&geom->friction[1]);
if (!newton_torsional || !newton_torsional.HasAuthoredValue()) {
mju_warning("Prim '%s' uses deprecated mjc:torsionalfriction. "
"Please migrate to newton:torsionalFriction.",
material_prim.GetPath().GetText());
}
} else if (newton_torsional && newton_torsional.HasAuthoredValue()) {
float val;
newton_torsional.Get(&val);
geom->friction[1] = val;
}
auto rolling_friction_attr = material_api.GetRollingFrictionAttr();
if (rolling_friction_attr.HasAuthoredValue()) {
rolling_friction_attr.Get(&geom->friction[2]);
// Rolling friction: same deprecation/fallback pattern.
auto mjc_rolling = material_api.GetRollingFrictionAttr();
auto newton_rolling = material_prim.GetAttribute(
kNewtonTokens->newtonRollingFriction);
if (mjc_rolling.HasAuthoredValue()) {
mjc_rolling.Get(&geom->friction[2]);
if (!newton_rolling || !newton_rolling.HasAuthoredValue()) {
mju_warning("Prim '%s' uses deprecated mjc:rollingfriction. "
"Please migrate to newton:rollingFriction.",
material_prim.GetPath().GetText());
}
} else if (newton_rolling && newton_rolling.HasAuthoredValue()) {
float val;
newton_rolling.Get(&val);
geom->friction[2] = val;
}
}
@@ -1789,11 +1925,13 @@ void ParseUsdPhysicsCollider(mjSpec* spec,
if (bound_material) {
pxr::UsdPrim bound_material_prim = bound_material.GetPrim();
if (bound_material_prim.HasAPI<pxr::UsdPhysicsMaterialAPI>() ||
bound_material_prim.HasAPI<pxr::MjcPhysicsMaterialAPI>()) {
bound_material_prim.HasAPI<pxr::MjcPhysicsMaterialAPI>() ||
bound_material_prim.HasAPI(kNewtonTokens->NewtonMaterialAPI)) {
ParseUsdPhysicsMaterialAPI(
geom, pxr::UsdPhysicsMaterialAPI(bound_material_prim));
ParseMjcPhysicsMaterialAPI(
geom, pxr::MjcPhysicsMaterialAPI(bound_material_prim));
geom, bound_material_prim,
pxr::MjcPhysicsMaterialAPI(bound_material_prim));
}
pxr::SdfPath material_path = bound_material_prim.GetPath();
mjsMaterial* material = nullptr;
@@ -1826,7 +1964,9 @@ void ParseUsdPhysicsCollider(mjSpec* spec,
if (!MaybeParseGeomPrimitive(prim, geom, caches.xform_cache)) {
mjsMesh* mesh = ParseUsdMesh(spec, prim, geom, caches.xform_cache);
if (mesh != nullptr && prim.HasAPI<pxr::MjcPhysicsMeshCollisionAPI>()) {
if (mesh != nullptr &&
(prim.HasAPI<pxr::MjcPhysicsMeshCollisionAPI>() ||
prim.HasAPI(kNewtonTokens->NewtonMeshCollisionAPI))) {
ParseMjcPhysicsMeshCollisionAPI(mesh,
pxr::MjcPhysicsMeshCollisionAPI(prim));
}
@@ -1877,8 +2017,8 @@ void ParseMjcEqualityAPISolverParams(
void ParseConstraint(mjSpec* spec, const pxr::UsdPrim& prim, mjsBody* body,
pxr::UsdGeomXformCache& xform_cache) {
if (prim.HasAPI<pxr::MjcPhysicsEqualityJointAPI>()) {
// Handle MjcPhysicsEqualityJointAPI on revolute/prismatic joints.
if (prim.HasAPI<pxr::MjcPhysicsEqualityJointAPI>() ||
prim.HasAPI(kNewtonTokens->NewtonMimicAPI)) {
pxr::MjcPhysicsEqualityJointAPI eq_joint_api(prim);
mjsEquality* eq = mjs_addEquality(spec, nullptr);
eq->type = mjEQ_JOINT;
@@ -1889,20 +2029,48 @@ void ParseConstraint(mjSpec* spec, const pxr::UsdPrim& prim, mjsBody* body,
eq->objtype = mjOBJ_JOINT;
mjs_setString(eq->name1, prim.GetPath().GetAsString().c_str());
// Get the target joint (joint2) from the MjcEqualityAPI target
// relationship.
pxr::MjcPhysicsEqualityAPI equality_api(prim);
pxr::UsdRelationship target_rel = equality_api.GetMjcTargetRel();
// Target joint: prefer newton:mimicJoint, fall back to deprecated mjc:target
pxr::SdfPathVector targets;
target_rel.GetTargets(&targets);
if (!targets.empty()) {
auto newton_mimic_rel = prim.GetRelationship(kNewtonTokens->newtonMimicJoint);
if (newton_mimic_rel && newton_mimic_rel.GetTargets(&targets) && !targets.empty()) {
mjs_setString(eq->name2, targets[0].GetAsString().c_str());
} else {
auto mjc_target_rel = prim.GetRelationship(MjcPhysicsTokens->mjcTarget);
if (mjc_target_rel && mjc_target_rel.GetTargets(&targets) && !targets.empty()) {
mjs_setString(eq->name2, targets[0].GetAsString().c_str());
mju_warning("Prim '%s' uses deprecated mjc:target. "
"Please migrate to newton:mimicJoint.",
prim.GetPath().GetText());
}
}
// If no target, name2 remains empty, meaning joint1 is fixed to a constant.
// Parse individual coefficient attributes for the quartic polynomial.
eq_joint_api.GetCoef0Attr().Get(&eq->data[0]);
eq_joint_api.GetCoef1Attr().Get(&eq->data[1]);
// Coefficients: prefer Newton, fall back to deprecated MJC
auto newton_coef0 = prim.GetAttribute(kNewtonTokens->newtonMimicCoef0);
auto newton_coef1 = prim.GetAttribute(kNewtonTokens->newtonMimicCoef1);
if (newton_coef0 && newton_coef0.HasAuthoredValue()) {
float val;
newton_coef0.Get(&val);
eq->data[0] = val;
} else {
eq_joint_api.GetCoef0Attr().Get(&eq->data[0]);
if (eq_joint_api.GetCoef0Attr().HasAuthoredValue()) {
mju_warning("Prim '%s' uses deprecated mjc:coef0. "
"Please migrate to newton:mimicCoef0.",
prim.GetPath().GetText());
}
}
if (newton_coef1 && newton_coef1.HasAuthoredValue()) {
float val;
newton_coef1.Get(&val);
eq->data[1] = val;
} else {
eq_joint_api.GetCoef1Attr().Get(&eq->data[1]);
if (eq_joint_api.GetCoef1Attr().HasAuthoredValue()) {
mju_warning("Prim '%s' uses deprecated mjc:coef1. "
"Please migrate to newton:mimicCoef1.",
prim.GetPath().GetText());
}
}
eq_joint_api.GetCoef2Attr().Get(&eq->data[2]);
eq_joint_api.GetCoef3Attr().Get(&eq->data[3]);
eq_joint_api.GetCoef4Attr().Get(&eq->data[4]);
@@ -1910,7 +2078,27 @@ void ParseConstraint(mjSpec* spec, const pxr::UsdPrim& prim, mjsBody* body,
pxr::UsdPhysicsJoint joint(prim);
ParseJointEnabled(eq, joint);
ParseMjcEqualityAPISolverParams(eq, equality_api, prim);
// Solver params are now inline on MjcEqualityJointAPI
auto solref_attr = prim.GetAttribute(MjcPhysicsTokens->mjcSolref);
if (solref_attr.HasAuthoredValue()) {
pxr::VtDoubleArray solref;
solref_attr.Get(&solref);
if (solref.size() == mjNREF) {
for (int i = 0; i < mjNREF; ++i) {
eq->solref[i] = solref[i];
}
}
}
auto solimp_attr = prim.GetAttribute(MjcPhysicsTokens->mjcSolimp);
if (solimp_attr.HasAuthoredValue()) {
pxr::VtDoubleArray solimp;
solimp_attr.Get(&solimp);
if (solimp.size() == mjNIMP) {
for (int i = 0; i < mjNIMP; ++i) {
eq->solimp[i] = solimp[i];
}
}
}
} else if (prim.IsA<pxr::UsdPhysicsFixedJoint>() ||
prim.IsA<pxr::UsdPhysicsSphericalJoint>()) {
// Handle fixed joints as weld constraints, spherical joints as connect constraints
python/mujoco/experimental/dear_imgui/dear_imgui.cc
+1233
View File
File diff suppressed because it is too large Load Diff
python/mujoco/experimental/dear_imgui/dear_imgui_macros.h
+326
View File
@@ -0,0 +1,326 @@
// Copyright 2026 DeepMind Technologies Limited
//
// 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
//
// https://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.
#ifndef MUJOCO_PYTHON_EXPERIMENTAL_DEAR_IMGUI_DEAR_IMGUI_MACROS_H_
#define MUJOCO_PYTHON_EXPERIMENTAL_DEAR_IMGUI_DEAR_IMGUI_MACROS_H_
// WARNING: This file is intended for internal use by dear_imgui libraries ONLY!
//
// The macros defined here use short, generic names (DEF0, ARG_ID, etc.) and
// are NOT #undef'd. Including this header elsewhere may cause naming conflicts.
//
// Define NAMESPACE to be the ImGui library you're binding before including this
// header, e.g. #define NAMESPACE ImGui
//
// ============================================================================
// Quick Reference
// ============================================================================
//
// DEFn(Name, Args...) // Binds NAMESPACE::Name as Name in Python
// DEFn_AS(CppName, PyName, Args...) // Binds NAMESPACE::CppName as PyName
// DEFn_F(PyName, Args..., { CppBody }) // Custom implementation
//
// Where 'n' is the number of arguments (0-9).
//
// ============================================================================
// Argument Format
// ============================================================================
//
// Each argument is a tuple: (Type, name, DefaultValue)
//
// - No default value: (ImString, label, ) // trailing comma required
// - With default value: (int, flags, = 0) // include the '='
// - Complex defaults: (const ImVec2&, size, = ImVec2_Zero)
//
// NOTE: Default values cannot contain commas. Use predefined constants like
// ImVec2_Zero, ImVec4_One, etc.
//
// ============================================================================
// Examples (from dear_imgui.cc)
// ============================================================================
//
// Simple binding - NAMESPACE::End() exposed as End():
// DEF0(End);
//
// Binding with arguments:
// DEF4(BeginChild,
// (ImString, str_id, ),
// (const ImVec2&, size, = ImVec2_Zero),
// (ImGuiChildFlags, child_flags, = 0),
// (ImGuiWindowFlags, window_flags, = 0));
//
// Overloaded function - NAMESPACE::BeginChild(ImGuiID) exposed as BeginChildId():
// DEF4_AS(BeginChild, BeginChildId,
// (ImGuiID, id, ),
// (const ImVec2&, size, = ImVec2_Zero),
// (ImGuiChildFlags, child_flags, = 0),
// (ImGuiWindowFlags, window_flags, = 0));
//
// Custom implementation using DEFn_F is needed when:
//
// 1. Variadic functions (e.g., Text, TextColored)
//
// C++ variadic functions (those with "...") cannot be bound directly
// because the type/count of arguments is unknown at compile time. Use a
// wrapper that calls the function with a fixed format. Also note that
// user-controlled format strings are a security risk (format string
// attacks). Always use "%s":
//
// DEF1_F(Text, (ImString, txt, ), {
// return NAMESPACE::Text("%s", txt);
// });
//
// 2. Output pointer parameters (e.g., Checkbox, SliderFloat)
//
// Python doesn't have output pointers, so return modified values as a
// tuple:
//
// DEF2_F(Checkbox, (ImString, label, ), (bool*, v, ), {
// auto result = NAMESPACE::Checkbox(label, v);
// return std::make_tuple(result, *v);
// });
//
// 3. Type conversions (e.g., Image, ImageButton)
//
// Some C++ types don't have Python equivalents. For example, ImTextureID
// is a void* (opaque pointer), which pybind11 can't automatically convert.
// Accept a Python-friendly type (like long) and cast it:
//
// DEF2_F(Image, (long, tex_id, ), (const ImVec2&, size, ), {
// return NAMESPACE::Image(reinterpret_cast<void*>(tex_id), size);
// });
// ============================================================================
// Internal helper macros (not intended to be called directly by binding code)
// ============================================================================
// Extracts the type and name of an argument tuple.
// Example: ARG_DECL((float, alpha, = 1.0f)) -> float alpha
#define ARG_DECL_X(T_, N_, V_) T_ N_
#define ARG_DECL(A_) ARG_DECL_X A_
// Extracts the identifier of an argument tuple.
// Example: ARG_ID((float, alpha, = 1.0f)) -> alpha
#define ARG_ID_X(T_, N_, V_) N_
#define ARG_ID(A_) ARG_ID_X A_
// Extracts the name of an argument tuple as a quoted string literal.
// Example: ARG_NAME((float, alpha, = 1.0f)) -> "alpha"
#define ARG_NAME_X(T_, N_, V_) #N_
#define ARG_NAME(A_) ARG_NAME_X A_
// Extracts the default value of an argument tuple.
// Example: ARG_DEFVAL((float, alpha, = 1.0f)) -> = 1.0f
#define ARG_DEFVAL_X(T_, N_, V_) V_
#define ARG_DEFVAL(A_) ARG_DEFVAL_X A_
// ============================================================================
// Public macros for binding code
// ============================================================================
//
#define DEF0_F(N, FN) \
m.def(#N, []( \
) FN \
);
#define DEF1_F(N, A1, FN) \
m.def(#N, []( \
ARG_DECL(A1) \
) FN, \
py::arg(ARG_NAME(A1)) ARG_DEFVAL(A1) \
);
#define DEF2_F(N, A1, A2, FN) \
m.def(#N, []( \
ARG_DECL(A1), \
ARG_DECL(A2) \
) FN, \
py::arg(ARG_NAME(A1)) ARG_DEFVAL(A1), \
py::arg(ARG_NAME(A2)) ARG_DEFVAL(A2) \
);
#define DEF3_F(N, A1, A2, A3, FN) \
m.def(#N, []( \
ARG_DECL(A1), \
ARG_DECL(A2), \
ARG_DECL(A3) \
) FN, \
py::arg(ARG_NAME(A1)) ARG_DEFVAL(A1), \
py::arg(ARG_NAME(A2)) ARG_DEFVAL(A2), \
py::arg(ARG_NAME(A3)) ARG_DEFVAL(A3) \
);
#define DEF4_F(N, A1, A2, A3, A4, FN) \
m.def(#N, []( \
ARG_DECL(A1), \
ARG_DECL(A2), \
ARG_DECL(A3), \
ARG_DECL(A4) \
) FN, \
py::arg(ARG_NAME(A1)) ARG_DEFVAL(A1), \
py::arg(ARG_NAME(A2)) ARG_DEFVAL(A2), \
py::arg(ARG_NAME(A3)) ARG_DEFVAL(A3), \
py::arg(ARG_NAME(A4)) ARG_DEFVAL(A4) \
);
#define DEF5_F(N, A1, A2, A3, A4, A5, FN) \
m.def(#N, []( \
ARG_DECL(A1), \
ARG_DECL(A2), \
ARG_DECL(A3), \
ARG_DECL(A4), \
ARG_DECL(A5) \
) FN, \
py::arg(ARG_NAME(A1)) ARG_DEFVAL(A1), \
py::arg(ARG_NAME(A2)) ARG_DEFVAL(A2), \
py::arg(ARG_NAME(A3)) ARG_DEFVAL(A3), \
py::arg(ARG_NAME(A4)) ARG_DEFVAL(A4), \
py::arg(ARG_NAME(A5)) ARG_DEFVAL(A5) \
);
#define DEF6_F(N, A1, A2, A3, A4, A5, A6, FN) \
m.def(#N, []( \
ARG_DECL(A1), \
ARG_DECL(A2), \
ARG_DECL(A3), \
ARG_DECL(A4), \
ARG_DECL(A5), \
ARG_DECL(A6) \
) FN, \
py::arg(ARG_NAME(A1)) ARG_DEFVAL(A1), \
py::arg(ARG_NAME(A2)) ARG_DEFVAL(A2), \
py::arg(ARG_NAME(A3)) ARG_DEFVAL(A3), \
py::arg(ARG_NAME(A4)) ARG_DEFVAL(A4), \
py::arg(ARG_NAME(A5)) ARG_DEFVAL(A5), \
py::arg(ARG_NAME(A6)) ARG_DEFVAL(A6) \
);
#define DEF7_F(N, A1, A2, A3, A4, A5, A6, A7, FN) \
m.def(#N, []( \
ARG_DECL(A1), \
ARG_DECL(A2), \
ARG_DECL(A3), \
ARG_DECL(A4), \
ARG_DECL(A5), \
ARG_DECL(A6), \
ARG_DECL(A7) \
) FN, \
py::arg(ARG_NAME(A1)) ARG_DEFVAL(A1), \
py::arg(ARG_NAME(A2)) ARG_DEFVAL(A2), \
py::arg(ARG_NAME(A3)) ARG_DEFVAL(A3), \
py::arg(ARG_NAME(A4)) ARG_DEFVAL(A4), \
py::arg(ARG_NAME(A5)) ARG_DEFVAL(A5), \
py::arg(ARG_NAME(A6)) ARG_DEFVAL(A6), \
py::arg(ARG_NAME(A7)) ARG_DEFVAL(A7) \
);
#define DEF8_F(N, A1, A2, A3, A4, A5, A6, A7, A8, FN) \
m.def(#N, []( \
ARG_DECL(A1), \
ARG_DECL(A2), \
ARG_DECL(A3), \
ARG_DECL(A4), \
ARG_DECL(A5), \
ARG_DECL(A6), \
ARG_DECL(A7), \
ARG_DECL(A8) \
) FN, \
py::arg(ARG_NAME(A1)) ARG_DEFVAL(A1), \
py::arg(ARG_NAME(A2)) ARG_DEFVAL(A2), \
py::arg(ARG_NAME(A3)) ARG_DEFVAL(A3), \
py::arg(ARG_NAME(A4)) ARG_DEFVAL(A4), \
py::arg(ARG_NAME(A5)) ARG_DEFVAL(A5), \
py::arg(ARG_NAME(A6)) ARG_DEFVAL(A6), \
py::arg(ARG_NAME(A7)) ARG_DEFVAL(A7), \
py::arg(ARG_NAME(A8)) ARG_DEFVAL(A8) \
);
#define DEF9_F(N, A1, A2, A3, A4, A5, A6, A7, A8, A9, FN) \
m.def(#N, []( \
ARG_DECL(A1), \
ARG_DECL(A2), \
ARG_DECL(A3), \
ARG_DECL(A4), \
ARG_DECL(A5), \
ARG_DECL(A6), \
ARG_DECL(A7), \
ARG_DECL(A8), \
ARG_DECL(A9) \
) FN, \
py::arg(ARG_NAME(A1)) ARG_DEFVAL(A1), \
py::arg(ARG_NAME(A2)) ARG_DEFVAL(A2), \
py::arg(ARG_NAME(A3)) ARG_DEFVAL(A3), \
py::arg(ARG_NAME(A4)) ARG_DEFVAL(A4), \
py::arg(ARG_NAME(A5)) ARG_DEFVAL(A5), \
py::arg(ARG_NAME(A6)) ARG_DEFVAL(A6), \
py::arg(ARG_NAME(A7)) ARG_DEFVAL(A7), \
py::arg(ARG_NAME(A8)) ARG_DEFVAL(A8), \
py::arg(ARG_NAME(A9)) ARG_DEFVAL(A9) \
);
#define DEF10_F(N, A1, A2, A3, A4, A5, A6, A7, A8, A9, A10, FN) \
m.def(#N, []( \
ARG_DECL(A1), \
ARG_DECL(A2), \
ARG_DECL(A3), \
ARG_DECL(A4), \
ARG_DECL(A5), \
ARG_DECL(A6), \
ARG_DECL(A7), \
ARG_DECL(A8), \
ARG_DECL(A9), \
ARG_DECL(A10) \
) FN, \
py::arg(ARG_NAME(A1)) ARG_DEFVAL(A1), \
py::arg(ARG_NAME(A2)) ARG_DEFVAL(A2), \
py::arg(ARG_NAME(A3)) ARG_DEFVAL(A3), \
py::arg(ARG_NAME(A4)) ARG_DEFVAL(A4), \
py::arg(ARG_NAME(A5)) ARG_DEFVAL(A5), \
py::arg(ARG_NAME(A6)) ARG_DEFVAL(A6), \
py::arg(ARG_NAME(A7)) ARG_DEFVAL(A7), \
py::arg(ARG_NAME(A8)) ARG_DEFVAL(A8), \
py::arg(ARG_NAME(A9)) ARG_DEFVAL(A9), \
py::arg(ARG_NAME(A10)) ARG_DEFVAL(A10) \
);
// NOLINTBEGIN(whitespace/line_length)
#define DEF0_AS(N, AS) DEF0_F(AS, { return NAMESPACE::N(); } )
#define DEF1_AS(N, AS, A1) DEF1_F(AS, A1, { return NAMESPACE::N(ARG_ID(A1)); } )
#define DEF2_AS(N, AS, A1, A2) DEF2_F(AS, A1, A2, { return NAMESPACE::N(ARG_ID(A1), ARG_ID(A2)); } )
#define DEF3_AS(N, AS, A1, A2, A3) DEF3_F(AS, A1, A2, A3, { return NAMESPACE::N(ARG_ID(A1), ARG_ID(A2), ARG_ID(A3)); } )
#define DEF4_AS(N, AS, A1, A2, A3, A4) DEF4_F(AS, A1, A2, A3, A4, { return NAMESPACE::N(ARG_ID(A1), ARG_ID(A2), ARG_ID(A3), ARG_ID(A4)); } )
#define DEF5_AS(N, AS, A1, A2, A3, A4, A5) DEF5_F(AS, A1, A2, A3, A4, A5, { return NAMESPACE::N(ARG_ID(A1), ARG_ID(A2), ARG_ID(A3), ARG_ID(A4), ARG_ID(A5)); } )
#define DEF6_AS(N, AS, A1, A2, A3, A4, A5, A6) DEF6_F(AS, A1, A2, A3, A4, A5, A6, { return NAMESPACE::N(ARG_ID(A1), ARG_ID(A2), ARG_ID(A3), ARG_ID(A4), ARG_ID(A5), ARG_ID(A6)); } )
#define DEF7_AS(N, AS, A1, A2, A3, A4, A5, A6, A7) DEF7_F(AS, A1, A2, A3, A4, A5, A6, A7, { return NAMESPACE::N(ARG_ID(A1), ARG_ID(A2), ARG_ID(A3), ARG_ID(A4), ARG_ID(A5), ARG_ID(A6), ARG_ID(A7)); } )
#define DEF8_AS(N, AS, A1, A2, A3, A4, A5, A6, A7, A8) DEF8_F(AS, A1, A2, A3, A4, A5, A6, A7, A8, { return NAMESPACE::N(ARG_ID(A1), ARG_ID(A2), ARG_ID(A3), ARG_ID(A4), ARG_ID(A5), ARG_ID(A6), ARG_ID(A7), ARG_ID(A8)); } )
#define DEF9_AS(N, AS, A1, A2, A3, A4, A5, A6, A7, A8, A9) DEF9_F(AS, A1, A2, A3, A4, A5, A6, A7, A8, A9, { return NAMESPACE::N(ARG_ID(A1), ARG_ID(A2), ARG_ID(A3), ARG_ID(A4), ARG_ID(A5), ARG_ID(A6), ARG_ID(A7), ARG_ID(A8), ARG_ID(A9)); } )
#define DEF10_AS(N, AS, A1, A2, A3, A4, A5, A6, A7, A8, A9, A10) DEF10_F(AS, A1, A2, A3, A4, A5, A6, A7, A8, A9, A10, { return NAMESPACE::N(ARG_ID(A1), ARG_ID(A2), ARG_ID(A3), ARG_ID(A4), ARG_ID(A5), ARG_ID(A6), ARG_ID(A7), ARG_ID(A8), ARG_ID(A9), ARG_ID(A10)); } )
#define DEF0(N) DEF0_F(N, { return NAMESPACE::N(); } )
#define DEF1(N, A1) DEF1_F(N, A1, { return NAMESPACE::N(ARG_ID(A1)); } )
#define DEF2(N, A1, A2) DEF2_F(N, A1, A2, { return NAMESPACE::N(ARG_ID(A1), ARG_ID(A2)); } )
#define DEF3(N, A1, A2, A3) DEF3_F(N, A1, A2, A3, { return NAMESPACE::N(ARG_ID(A1), ARG_ID(A2), ARG_ID(A3)); } )
#define DEF4(N, A1, A2, A3, A4) DEF4_F(N, A1, A2, A3, A4, { return NAMESPACE::N(ARG_ID(A1), ARG_ID(A2), ARG_ID(A3), ARG_ID(A4)); } )
#define DEF5(N, A1, A2, A3, A4, A5) DEF5_F(N, A1, A2, A3, A4, A5, { return NAMESPACE::N(ARG_ID(A1), ARG_ID(A2), ARG_ID(A3), ARG_ID(A4), ARG_ID(A5)); } )
#define DEF6(N, A1, A2, A3, A4, A5, A6) DEF6_F(N, A1, A2, A3, A4, A5, A6, { return NAMESPACE::N(ARG_ID(A1), ARG_ID(A2), ARG_ID(A3), ARG_ID(A4), ARG_ID(A5), ARG_ID(A6)); } )
#define DEF7(N, A1, A2, A3, A4, A5, A6, A7) DEF7_F(N, A1, A2, A3, A4, A5, A6, A7, { return NAMESPACE::N(ARG_ID(A1), ARG_ID(A2), ARG_ID(A3), ARG_ID(A4), ARG_ID(A5), ARG_ID(A6), ARG_ID(A7)); } )
#define DEF8(N, A1, A2, A3, A4, A5, A6, A7, A8) DEF8_F(N, A1, A2, A3, A4, A5, A6, A7, A8, { return NAMESPACE::N(ARG_ID(A1), ARG_ID(A2), ARG_ID(A3), ARG_ID(A4), ARG_ID(A5), ARG_ID(A6), ARG_ID(A7), ARG_ID(A8)); } )
#define DEF9(N, A1, A2, A3, A4, A5, A6, A7, A8, A9) DEF9_F(N, A1, A2, A3, A4, A5, A6, A7, A8, A9, { return NAMESPACE::N(ARG_ID(A1), ARG_ID(A2), ARG_ID(A3), ARG_ID(A4), ARG_ID(A5), ARG_ID(A6), ARG_ID(A7), ARG_ID(A8), ARG_ID(A9)); } )
#define DEF10(N, A1, A2, A3, A4, A5, A6, A7, A8, A9, A10) DEF10_F(N, A1, A2, A3, A4, A5, A6, A7, A8, A9, A10, { return NAMESPACE::N(ARG_ID(A1), ARG_ID(A2), ARG_ID(A3), ARG_ID(A4), ARG_ID(A5), ARG_ID(A6), ARG_ID(A7), ARG_ID(A8), ARG_ID(A9), ARG_ID(A10)); } )
// NOLINTEND(whitespace/line_length)
#endif // MUJOCO_PYTHON_EXPERIMENTAL_DEAR_IMGUI_DEAR_IMGUI_MACROS_H_
python/mujoco/experimental/implot/implot.cc
+438
View File
@@ -0,0 +1,438 @@
// Copyright 2026 DeepMind Technologies Limited
//
// 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
//
// https://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.
#define NAMESPACE ImPlot
#include "dear_imgui_macros.h"
#include <implot.h>
#include <pybind11/eval.h>
#include <pybind11/pybind11.h>
#include <pybind11/pytypes.h>
#include <pybind11/stl.h>
// NOLINTBEGIN(whitespace/line_length)
namespace py = pybind11;
using ImString = const char*;
static constexpr const ImVec2 ImVec2_Zero = ImVec2(0.0f, 0.0f);
static constexpr const ImVec2 ImVec2_One = ImVec2(1.0f, 1.0f);
static constexpr const ImVec2 ImVec2_NegOne_Zero = ImVec2(-1.0f, 0.0f);
static constexpr const ImVec4 ImVec4_Zero = ImVec4(0.0f, 0.0f, 0.0f, 0.0f);
static constexpr const ImVec4 ImVec4_One = ImVec4(1.0f, 1.0f, 1.0f, 1.0f);
static constexpr const ImPlotRect ImPlotRect_Default{};
static constexpr const ImPlotRange ImPlotRange_Default{};
PYBIND11_MODULE(implot, m) {
// Import dear_imgui to make types like ImVec2 available.
py::module_::import("mujoco.experimental.dear_imgui.dear_imgui");
// Types.
py::class_<ImPlotPoint>(m, "Point")
.def(py::init<>())
.def(py::init<double, double>(), py::arg("_x"), py::arg("_y"))
.def_readwrite("x", &ImPlotPoint::x)
.def_readwrite("y", &ImPlotPoint::y);
py::class_<ImPlotRange>(m, "Range")
.def(py::init<>())
.def(py::init<double, double>(), py::arg("_min"), py::arg("_max"))
.def_readwrite("min", &ImPlotRange::Min)
.def_readwrite("max", &ImPlotRange::Max);
// py::class_<ImPlotRect>(m, "Rect")
// .def(py::init<>())
// .def(py::init<ImPlotRange, ImPlotRange>(), py::arg("_x"), py::arg("_y"))
// .def_readwrite("x", &ImPlotRect::X)
// .def_readwrite("y", &ImPlotRect::Y);
// Enumerations.
py::enum_<ImAxis_>(m, "Axis")
.value("X1", ImAxis_X1)
.value("X2", ImAxis_X2)
.value("X3", ImAxis_X3)
.value("Y1", ImAxis_Y1)
.value("Y2", ImAxis_Y2)
.value("Y3", ImAxis_Y3);
py::enum_<ImPlotFlags_>(m, "Flags")
.value("None", ImPlotFlags_None)
.value("NoTitle", ImPlotFlags_NoTitle)
.value("NoLegend", ImPlotFlags_NoLegend)
.value("NoMouseText", ImPlotFlags_NoMouseText)
.value("NoInputs", ImPlotFlags_NoInputs)
.value("NoMenus", ImPlotFlags_NoMenus)
.value("NoBoxSelect", ImPlotFlags_NoBoxSelect)
.value("NoFrame", ImPlotFlags_NoFrame)
.value("Equal", ImPlotFlags_Equal)
.value("Crosshairs", ImPlotFlags_Crosshairs)
.value("CanvasOnly", ImPlotFlags_CanvasOnly);
py::enum_<ImPlotAxisFlags_>(m, "AxisFlags")
.value("None", ImPlotAxisFlags_None)
.value("NoLabel", ImPlotAxisFlags_NoLabel)
.value("NoGridLines", ImPlotAxisFlags_NoGridLines)
.value("NoTickMarks", ImPlotAxisFlags_NoTickMarks)
.value("NoTickLabels", ImPlotAxisFlags_NoTickLabels)
.value("NoInitialFit", ImPlotAxisFlags_NoInitialFit)
.value("NoMenus", ImPlotAxisFlags_NoMenus)
.value("NoSideSwitch", ImPlotAxisFlags_NoSideSwitch)
.value("NoHighlight", ImPlotAxisFlags_NoHighlight)
.value("Opposite", ImPlotAxisFlags_Opposite)
.value("Foreground", ImPlotAxisFlags_Foreground)
.value("Invert", ImPlotAxisFlags_Invert)
.value("AutoFit", ImPlotAxisFlags_AutoFit)
.value("RangeFit", ImPlotAxisFlags_RangeFit)
.value("PanStretch", ImPlotAxisFlags_PanStretch)
.value("LockMin", ImPlotAxisFlags_LockMin)
.value("LockMax", ImPlotAxisFlags_LockMax)
.value("Lock", ImPlotAxisFlags_Lock)
.value("NoDecorations", ImPlotAxisFlags_NoDecorations)
.value("AuxDefault", ImPlotAxisFlags_AuxDefault);
py::enum_<ImPlotSubplotFlags_>(m, "SubplotFlags")
.value("None", ImPlotSubplotFlags_None)
.value("NoTitle", ImPlotSubplotFlags_NoTitle)
.value("NoLegend", ImPlotSubplotFlags_NoLegend)
.value("NoMenus", ImPlotSubplotFlags_NoMenus)
.value("NoResize", ImPlotSubplotFlags_NoResize)
.value("NoAlign", ImPlotSubplotFlags_NoAlign)
.value("ShareItems", ImPlotSubplotFlags_ShareItems)
.value("LinkRows", ImPlotSubplotFlags_LinkRows)
.value("LinkCols", ImPlotSubplotFlags_LinkCols)
.value("LinkAllX", ImPlotSubplotFlags_LinkAllX)
.value("LinkAllY", ImPlotSubplotFlags_LinkAllY)
.value("ColMajor", ImPlotSubplotFlags_ColMajor);
py::enum_<ImPlotLegendFlags_>(m, "LegendFlags")
.value("None", ImPlotLegendFlags_None)
.value("NoButtons", ImPlotLegendFlags_NoButtons)
.value("NoHighlightItem", ImPlotLegendFlags_NoHighlightItem)
.value("NoHighlightAxis", ImPlotLegendFlags_NoHighlightAxis)
.value("NoMenus", ImPlotLegendFlags_NoMenus)
.value("Outside", ImPlotLegendFlags_Outside)
.value("Horizontal", ImPlotLegendFlags_Horizontal)
.value("Sort", ImPlotLegendFlags_Sort)
.value("Reverse", ImPlotLegendFlags_Reverse);
py::enum_<ImPlotMouseTextFlags_>(m, "MouseTextFlags")
.value("None", ImPlotMouseTextFlags_None)
.value("NoAuxAxes", ImPlotMouseTextFlags_NoAuxAxes)
.value("NoFormat", ImPlotMouseTextFlags_NoFormat)
.value("ShowAlways", ImPlotMouseTextFlags_ShowAlways);
py::enum_<ImPlotDragToolFlags_>(m, "DragToolFlags")
.value("None", ImPlotDragToolFlags_None)
.value("NoCursors", ImPlotDragToolFlags_NoCursors)
.value("NoFit", ImPlotDragToolFlags_NoFit)
.value("NoInputs", ImPlotDragToolFlags_NoInputs)
.value("Delayed", ImPlotDragToolFlags_Delayed);
py::enum_<ImPlotColormapScaleFlags_>(m, "ColormapScaleFlags")
.value("None", ImPlotColormapScaleFlags_None)
.value("NoLabel", ImPlotColormapScaleFlags_NoLabel)
.value("Opposite", ImPlotColormapScaleFlags_Opposite)
.value("Invert", ImPlotColormapScaleFlags_Invert);
py::enum_<ImPlotItemFlags_>(m, "ItemFlags")
.value("None", ImPlotItemFlags_None)
.value("NoLegend", ImPlotItemFlags_NoLegend)
.value("NoFit", ImPlotItemFlags_NoFit);
py::enum_<ImPlotLineFlags_>(m, "LineFlags")
.value("None", ImPlotLineFlags_None)
.value("Segments", ImPlotLineFlags_Segments)
.value("Loop", ImPlotLineFlags_Loop)
.value("SkipNaN", ImPlotLineFlags_SkipNaN)
.value("NoClip", ImPlotLineFlags_NoClip)
.value("Shaded", ImPlotLineFlags_Shaded);
py::enum_<ImPlotScatterFlags_>(m, "ScatterFlags")
.value("None", ImPlotScatterFlags_None)
.value("NoClip", ImPlotScatterFlags_NoClip);
py::enum_<ImPlotStairsFlags_>(m, "StairsFlags")
.value("None", ImPlotStairsFlags_None)
.value("PreStep", ImPlotStairsFlags_PreStep)
.value("Shaded", ImPlotStairsFlags_Shaded);
py::enum_<ImPlotShadedFlags_>(m, "ShadedFlags")
.value("None", ImPlotShadedFlags_None);
py::enum_<ImPlotBarsFlags_>(m, "BarsFlags")
.value("None", ImPlotBarsFlags_None)
.value("Horizontal", ImPlotBarsFlags_Horizontal);
py::enum_<ImPlotBarGroupsFlags_>(m, "BarGroupsFlags")
.value("None", ImPlotBarGroupsFlags_None)
.value("Horizontal", ImPlotBarGroupsFlags_Horizontal)
.value("Stacked", ImPlotBarGroupsFlags_Stacked);
py::enum_<ImPlotErrorBarsFlags_>(m, "ErrorBarsFlags")
.value("None", ImPlotErrorBarsFlags_None)
.value("Horizontal", ImPlotErrorBarsFlags_Horizontal);
py::enum_<ImPlotStemsFlags_>(m, "StemsFlags")
.value("None", ImPlotStemsFlags_None)
.value("Horizontal", ImPlotStemsFlags_Horizontal);
py::enum_<ImPlotInfLinesFlags_>(m, "InfLinesFlags")
.value("None", ImPlotInfLinesFlags_None)
.value("Horizontal", ImPlotInfLinesFlags_Horizontal);
py::enum_<ImPlotPieChartFlags_>(m, "PieChartFlags")
.value("None", ImPlotPieChartFlags_None)
.value("Normalize", ImPlotPieChartFlags_Normalize)
.value("IgnoreHidden", ImPlotPieChartFlags_IgnoreHidden)
.value("Exploding", ImPlotPieChartFlags_Exploding);
py::enum_<ImPlotHeatmapFlags_>(m, "HeatmapFlags")
.value("None", ImPlotHeatmapFlags_None)
.value("ColMajor", ImPlotHeatmapFlags_ColMajor);
py::enum_<ImPlotHistogramFlags_>(m, "HistogramFlags")
.value("None", ImPlotHistogramFlags_None)
.value("Horizontal", ImPlotHistogramFlags_Horizontal)
.value("Cumulative", ImPlotHistogramFlags_Cumulative)
.value("Density", ImPlotHistogramFlags_Density)
.value("NoOutliers", ImPlotHistogramFlags_NoOutliers)
.value("ColMajor", ImPlotHistogramFlags_ColMajor);
py::enum_<ImPlotDigitalFlags_>(m, "DigitalFlags")
.value("ImPlotNone", ImPlotDigitalFlags_None);
py::enum_<ImPlotImageFlags_>(m, "ImageFlags")
.value("None", ImPlotImageFlags_None);
py::enum_<ImPlotTextFlags_>(m, "TextFlags")
.value("None", ImPlotTextFlags_None)
.value("Vertical", ImPlotTextFlags_Vertical);
py::enum_<ImPlotDummyFlags_>(m, "DummyFlags")
.value("None", ImPlotDummyFlags_None);
py::enum_<ImPlotCond_>(m, "Cond")
.value("None", ImPlotCond_None)
.value("Always", ImPlotCond_Always)
.value("Once", ImPlotCond_Once);
py::enum_<ImPlotCol_>(m, "Col")
.value("Line", ImPlotCol_Line)
.value("Fill", ImPlotCol_Fill)
.value("MarkerOutline", ImPlotCol_MarkerOutline)
.value("MarkerFill", ImPlotCol_MarkerFill)
.value("ErrorBar", ImPlotCol_ErrorBar)
.value("FrameBg", ImPlotCol_FrameBg)
.value("PlotBg", ImPlotCol_PlotBg)
.value("PlotBorder", ImPlotCol_PlotBorder)
.value("LegendBg", ImPlotCol_LegendBg)
.value("LegendBorder", ImPlotCol_LegendBorder)
.value("LegendText", ImPlotCol_LegendText)
.value("TitleText", ImPlotCol_TitleText)
.value("InlayText", ImPlotCol_InlayText)
.value("AxisText", ImPlotCol_AxisText)
.value("AxisGrid", ImPlotCol_AxisGrid)
.value("AxisTick", ImPlotCol_AxisTick)
.value("AxisBg", ImPlotCol_AxisBg)
.value("AxisBgHovered", ImPlotCol_AxisBgHovered)
.value("AxisBgActive", ImPlotCol_AxisBgActive)
.value("Selection", ImPlotCol_Selection)
.value("Crosshairs", ImPlotCol_Crosshairs);
py::enum_<ImPlotStyleVar_>(m, "StyleVar")
.value("LineWeight", ImPlotStyleVar_LineWeight)
.value("Marker", ImPlotStyleVar_Marker)
.value("MarkerSize", ImPlotStyleVar_MarkerSize)
.value("MarkerWeight", ImPlotStyleVar_MarkerWeight)
.value("FillAlpha", ImPlotStyleVar_FillAlpha)
.value("ErrorBarSize", ImPlotStyleVar_ErrorBarSize)
.value("ErrorBarWeight", ImPlotStyleVar_ErrorBarWeight)
.value("DigitalBitHeight", ImPlotStyleVar_DigitalBitHeight)
.value("DigitalBitGap", ImPlotStyleVar_DigitalBitGap)
.value("PlotBorderSize", ImPlotStyleVar_PlotBorderSize)
.value("MinorAlpha", ImPlotStyleVar_MinorAlpha)
.value("MajorTickLen", ImPlotStyleVar_MajorTickLen)
.value("MinorTickLen", ImPlotStyleVar_MinorTickLen)
.value("MajorTickSize", ImPlotStyleVar_MajorTickSize)
.value("MinorTickSize", ImPlotStyleVar_MinorTickSize)
.value("MajorGridSize", ImPlotStyleVar_MajorGridSize)
.value("MinorGridSize", ImPlotStyleVar_MinorGridSize)
.value("PlotPadding", ImPlotStyleVar_PlotPadding)
.value("LabelPadding", ImPlotStyleVar_LabelPadding)
.value("LegendPadding", ImPlotStyleVar_LegendPadding)
.value("LegendInnerPadding", ImPlotStyleVar_LegendInnerPadding)
.value("LegendSpacing", ImPlotStyleVar_LegendSpacing)
.value("MousePosPadding", ImPlotStyleVar_MousePosPadding)
.value("AnnotationPadding", ImPlotStyleVar_AnnotationPadding)
.value("FitPadding", ImPlotStyleVar_FitPadding)
.value("PlotDefaultSize", ImPlotStyleVar_PlotDefaultSize)
.value("PlotMinSize", ImPlotStyleVar_PlotMinSize);
py::enum_<ImPlotScale_>(m, "Scale")
.value("ImPlotScale_Linear", ImPlotScale_Linear)
.value("ImPlotScale_Time", ImPlotScale_Time)
.value("ImPlotScale_Log10", ImPlotScale_Log10)
.value("ImPlotScale_SymLog", ImPlotScale_SymLog);
py::enum_<ImPlotMarker_>(m, "Marker")
.value("None", ImPlotMarker_None)
.value("Circle", ImPlotMarker_Circle)
.value("Square", ImPlotMarker_Square)
.value("Diamond", ImPlotMarker_Diamond)
.value("Up", ImPlotMarker_Up)
.value("Down", ImPlotMarker_Down)
.value("Left", ImPlotMarker_Left)
.value("Right", ImPlotMarker_Right)
.value("Cross", ImPlotMarker_Cross)
.value("Plus", ImPlotMarker_Plus)
.value("Asterisk", ImPlotMarker_Asterisk);
py::enum_<ImPlotLocation_>(m, "Location")
.value("Center", ImPlotLocation_Center)
.value("North", ImPlotLocation_North)
.value("South", ImPlotLocation_South)
.value("West", ImPlotLocation_West)
.value("East", ImPlotLocation_East)
.value("NorthWest", ImPlotLocation_NorthWest)
.value("NorthEast", ImPlotLocation_NorthEast)
.value("SouthWest", ImPlotLocation_SouthWest)
.value("SouthEast", ImPlotLocation_SouthEast);
// Functions.
DEF3(BeginPlot, (ImString, title_id, ), (const ImVec2&, size, = ImVec2_NegOne_Zero), (ImPlotFlags, flags, = 0));
DEF0(EndPlot);
DEF7(BeginSubplots, (ImString, title_id, ), (int, rows, ), (int, cols, ), (const ImVec2&, size, ), (ImPlotSubplotFlags, flags, = 0), (float*, row_ratios, = nullptr), (float*, col_ratios, = nullptr));
DEF0(EndSubplots);
DEF3(SetupAxis, (ImAxis, axis, ), (ImString, label, = nullptr), (ImPlotAxisFlags, flags, = 0));
DEF4(SetupAxisLimits, (ImAxis, axis, ), (double, v_min, ), (double, v_max, ), (ImPlotCond, cond, = ImPlotCond_Once));
DEF4_F(SetupAxisTicks, (ImAxis, axis, ), (std::vector<double>, values, ), (std::vector<std::string>, labels, ), (bool, keep_default, = false), {
std::vector<const char*> c_labels;
c_labels.reserve(labels.size());
for (const auto& l : labels) {
c_labels.push_back(l.c_str());
}
ImPlot::SetupAxisTicks(axis, values.data(), values.size(), c_labels.empty() ? nullptr : c_labels.data(), keep_default);
});
DEF3(SetupAxisLinks, (ImAxis, axis, ), (double*, link_min, ), (double*, link_max, ));
DEF2(SetupAxisFormat, (ImAxis, axis, ), (ImString, fmt, ));
DEF2(SetupAxisScale, (ImAxis, axis, ), (ImPlotScale, scale, ));
DEF3(SetupAxisLimitsConstraints, (ImAxis, axis, ), (double, v_min, ), (double, v_max, ));
DEF3(SetupAxisZoomConstraints, (ImAxis, axis, ), (double, z_min, ), (double, z_max, ));
DEF4(SetupAxes, (ImString, x_label, ), (ImString, y_label, ), (ImPlotAxisFlags, x_flags, = 0), (ImPlotAxisFlags, y_flags, = 0));
DEF5(SetupAxesLimits, (double, x_min, ), (double, x_max, ), (double, y_min, ), (double, y_max, ), (ImPlotCond, cond, = ImPlotCond_Once));
DEF2(SetupLegend, (ImPlotLocation, location, ), (ImPlotLegendFlags, flags, = 0));
DEF2(SetupMouseText, (ImPlotLocation, location, ), (ImPlotMouseTextFlags, flags, = 0));
DEF0(SetupFinish);
DEF4(SetNextAxisLimits, (ImAxis, axis, ), (double, v_min, ), (double, v_max, ), (ImPlotCond, cond, = ImPlotCond_Once));
DEF3(SetNextAxisLinks, (ImAxis, axis, ), (double*, link_min, ), (double*, link_max, ));
DEF1(SetNextAxisToFit, (ImAxis, axis, ));
DEF5(SetNextAxesLimits, (double, x_min, ), (double, x_max, ), (double, y_min, ), (double, y_max, ), (ImPlotCond, cond, = ImPlotCond_Once));
DEF0_F(SetNextAxesToFit, {
return ImPlot::SetNextAxesToFit();
});
DEF6_F(PlotLine, (ImString, label_id, ), (std::vector<double>, xs, ), (std::vector<double>, ys, ), (ImPlotLineFlags, flags, = 0), (int, offset, = 0), (int, stride, = sizeof(double)), {
return ImPlot::PlotLine(label_id, xs.data(), ys.data(), xs.size(), flags, offset, stride);
});
DEF6_F(PlotScatter, (ImString, label_id, ), (std::vector<double>, xs, ), (std::vector<double>, ys, ), (ImPlotScatterFlags, flags, = 0), (int, offset, = 0), (int, stride, = sizeof(double)), {
return ImPlot::PlotScatter(label_id, xs.data(), ys.data(), xs.size(), flags, offset, stride);
});
DEF6_F(PlotStairs, (ImString, label_id, ), (std::vector<double>, xs, ), (std::vector<double>, ys, ), (ImPlotStairsFlags, flags, = 0), (int, offset, = 0), (int, stride, = sizeof(double)), {
return ImPlot::PlotStairs(label_id, xs.data(), ys.data(), xs.size(), flags, offset, stride);
});
DEF7_F(PlotShaded, (ImString, label_id, ), (std::vector<double>, xs, ), (std::vector<double>, ys, ), (double, yref, = 0), (ImPlotShadedFlags, flags, = 0), (int, offset, = 0), (int, stride, = sizeof(double)), {
return ImPlot::PlotShaded(label_id, xs.data(), ys.data(), xs.size(), yref, flags, offset, stride);
});
DEF7_F(PlotShaded, (ImString, label_id, ), (std::vector<double>, xs, ), (std::vector<double>, ys1, ), (std::vector<double>, ys2, ), (ImPlotShadedFlags, flags, = 0), (int, offset, = 0), (int, stride, = sizeof(double)), {
return ImPlot::PlotShaded(label_id, xs.data(), ys1.data(), ys2.data(), xs.size(), flags, offset, stride);
});
DEF7_F(PlotBars, (ImString, label_id, ), (std::vector<double>, xs, ), (std::vector<double>, ys, ), (double, bar_size, ), (ImPlotBarsFlags, flags, = 0), (int, offset, = 0), (int, stride, = sizeof(double)), {
return ImPlot::PlotBars(label_id, xs.data(), ys.data(), xs.size(), bar_size, flags, offset, stride);
});
DEF7_F(PlotErrorBars, (ImString, label_id, ), (std::vector<double>, xs, ), (std::vector<double>, ys, ), (std::vector<double>, err, ), (ImPlotErrorBarsFlags, flags, = 0), (int, offset, = 0), (int, stride, = sizeof(double)), {
return ImPlot::PlotErrorBars(label_id, xs.data(), ys.data(), err.data(), xs.size(), flags, offset, stride);
});
DEF8_F(PlotErrorBars, (ImString, label_id, ), (std::vector<double>, xs, ), (std::vector<double>, ys, ), (std::vector<double>, neg, ), (std::vector<double>, pos, ), (ImPlotErrorBarsFlags, flags, = 0), (int, offset, = 0), (int, stride, = sizeof(double)), {
return ImPlot::PlotErrorBars(label_id, xs.data(), ys.data(), neg.data(), pos.data(), xs.size(), flags, offset, stride);
});
DEF7_F(PlotStems, (ImString, label_id, ), (std::vector<double>, xs, ), (std::vector<double>, ys, ), (double, ref, = 0), (ImPlotStemsFlags, flags, = 0), (int, offset, = 0), (int, stride, = sizeof(double)), {
return ImPlot::PlotStems(label_id, xs.data(), ys.data(), xs.size(), ref, flags, offset, stride);
});
DEF6_F(PlotDigital, (ImString, label_id, ), (std::vector<double>, xs, ), (std::vector<double>, ys, ), (ImPlotDigitalFlags, flags, = 0), (int, offset, = 0), (int, stride, = sizeof(double)), {
return ImPlot::PlotDigital(label_id, xs.data(), ys.data(), xs.size(), flags, offset, stride);
});
DEF8(PlotImage, (ImString, label_id, ), (ImTextureRef, tex_ref, ), (const ImPlotPoint&, bounds_min, ), (const ImPlotPoint&, bounds_max, ), (const ImVec2&, uv0, = ImVec2_Zero), (const ImVec2&, uv1, = ImVec2_One), (const ImVec4&, tint_col, = ImVec4_One), (ImPlotImageFlags, flags, = 0));
DEF5(PlotText, (ImString, text, ), (double, x, ), (double, y, ), (const ImVec2&, pix_offset, = ImVec2_Zero), (ImPlotTextFlags, flags, = 0));
DEF2(PlotDummy, (ImString, label_id, ), (ImPlotDummyFlags, flags, = 0));
DEF9(DragPoint, (int, id, ), (double*, x, ), (double*, y, ), (const ImVec4&, col, ), (float, size, = 4), (ImPlotDragToolFlags, flags, = 0), (bool*, out_clicked, = nullptr), (bool*, out_hovered, = nullptr), (bool*, out_held, = nullptr));
DEF8(DragLineX, (int, id, ), (double*, x, ), (const ImVec4&, col, ), (float, thickness, = 1), (ImPlotDragToolFlags, flags, = 0), (bool*, out_clicked, = nullptr), (bool*, out_hovered, = nullptr), (bool*, out_held, = nullptr));
DEF8(DragLineY, (int, id, ), (double*, y, ), (const ImVec4&, col, ), (float, thickness, = 1), (ImPlotDragToolFlags, flags, = 0), (bool*, out_clicked, = nullptr), (bool*, out_hovered, = nullptr), (bool*, out_held, = nullptr));
DEF10(DragRect, (int, id, ), (double*, x1, ), (double*, y1, ), (double*, x2, ), (double*, y2, ), (const ImVec4&, col, ), (ImPlotDragToolFlags, flags, = 0), (bool*, out_clicked, = nullptr), (bool*, out_hovered, = nullptr), (bool*, out_held, = nullptr));
DEF6(Annotation, (double, x, ), (double, y, ), (const ImVec4&, col, ), (const ImVec2&, pix_offset, ), (bool, clamp, ), (bool, round, = false));
DEF6_F(Annotation, (double, x, ), (double, y, ), (const ImVec4&, col, ), (const ImVec2&, pix_offset, ), (bool, clamp, ), (ImString, txt, ), {
return ImPlot::Annotation(x, y, col, pix_offset, clamp, "%s", txt);
});
DEF3(TagX, (double, x, ), (const ImVec4&, col, ), (bool, round, = false));
DEF3_F(TagX, (double, x, ), (const ImVec4&, col, ), (ImString, txt, ), {
return ImPlot::TagX(x, col, "%s", txt);
});
DEF3(TagY, (double, y, ), (const ImVec4&, col, ), (bool, round, = false));
DEF3_F(TagY, (double, y, ), (const ImVec4&, col, ), (ImString, txt, ), {
return ImPlot::TagY(y, col, "%s", txt);
});
DEF1(SetAxis, (ImAxis, axis, ));
DEF2(SetAxes, (ImAxis, x_axis, ), (ImAxis, y_axis, ));
DEF3(PixelsToPlot, (const ImVec2&, pix, ), (ImAxis, x_axis, = IMPLOT_AUTO), (ImAxis, y_axis, = IMPLOT_AUTO));
DEF4(PixelsToPlot, (float, x, ), (float, y, ), (ImAxis, x_axis, = IMPLOT_AUTO), (ImAxis, y_axis, = IMPLOT_AUTO));
DEF3(PlotToPixels, (const ImPlotPoint&, plt, ), (ImAxis, x_axis, = IMPLOT_AUTO), (ImAxis, y_axis, = IMPLOT_AUTO));
DEF4(PlotToPixels, (double, x, ), (double, y, ), (ImAxis, x_axis, = IMPLOT_AUTO), (ImAxis, y_axis, = IMPLOT_AUTO));
DEF0(GetPlotPos);
DEF0(GetPlotSize);
DEF2(GetPlotMousePos, (ImAxis, x_axis, = IMPLOT_AUTO), (ImAxis, y_axis, = IMPLOT_AUTO));
DEF2(GetPlotLimits, (ImAxis, x_axis, = IMPLOT_AUTO), (ImAxis, y_axis, = IMPLOT_AUTO));
DEF0(IsPlotHovered);
DEF1(IsAxisHovered, (ImAxis, axis, ));
DEF0(IsSubplotsHovered);
DEF0(IsPlotSelected);
DEF2(GetPlotSelection, (ImAxis, x_axis, = IMPLOT_AUTO), (ImAxis, y_axis, = IMPLOT_AUTO));
DEF0(CancelPlotSelection);
DEF2(HideNextItem, (bool, hidden, = true), (ImPlotCond, cond, = ImPlotCond_Once));
DEF2(BeginAlignedPlots, (ImString, group_id, ), (bool, vertical, = true));
DEF0(EndAlignedPlots);
DEF2(BeginLegendPopup, (ImString, label_id, ), (ImGuiMouseButton, mouse_button, = 1));
DEF0(EndLegendPopup);
DEF1(IsLegendEntryHovered, (ImString, label_id, ));
DEF0(BeginDragDropTargetPlot);
DEF1(BeginDragDropTargetAxis, (ImAxis, axis, ));
DEF0(BeginDragDropTargetLegend);
DEF0(EndDragDropTarget);
DEF1(BeginDragDropSourcePlot, (ImGuiDragDropFlags, flags, = 0));
DEF2(BeginDragDropSourceAxis, (ImAxis, axis, ), (ImGuiDragDropFlags, flags, = 0));
DEF2(BeginDragDropSourceItem, (ImString, label_id, ), (ImGuiDragDropFlags, flags, = 0));
DEF0(EndDragDropSource);
DEF2(PushStyleColor, (ImPlotCol, idx, ), (ImU32, col, ));
DEF2(PushStyleColor, (ImPlotCol, idx, ), (const ImVec4&, col, ));
DEF1(PopStyleColor, (int, count, = 1));
DEF2(PushStyleVar, (ImPlotStyleVar, idx, ), (float, val, ));
DEF2(PushStyleVar, (ImPlotStyleVar, idx, ), (int, val, ));
DEF2(PushStyleVar, (ImPlotStyleVar, idx, ), (const ImVec2&, val, ));
DEF1(PopStyleVar, (int, count, = 1));
DEF2(SetNextLineStyle, (const ImVec4&, col, = IMPLOT_AUTO_COL), (float, weight, = IMPLOT_AUTO));
DEF2(SetNextFillStyle, (const ImVec4&, col, = IMPLOT_AUTO_COL), (float, alpha_mod, = IMPLOT_AUTO));
DEF5(SetNextMarkerStyle, (ImPlotMarker, marker, = IMPLOT_AUTO), (float, size, = IMPLOT_AUTO), (const ImVec4&, fill, = IMPLOT_AUTO_COL), (float, weight, = IMPLOT_AUTO), (const ImVec4&, outline, = IMPLOT_AUTO_COL));
DEF3(SetNextErrorBarStyle, (const ImVec4&, col, = IMPLOT_AUTO_COL), (float, size, = IMPLOT_AUTO), (float, weight, = IMPLOT_AUTO));
DEF1(PushPlotClipRect, (float, expand, = 0));
DEF0(PopPlotClipRect);
}
// NOLINTEND(whitespace/line_length)
python/mujoco/experimental/studio/native_viewer.cc
+189
View File
@@ -0,0 +1,189 @@
// Copyright 2026 DeepMind Technologies Limited
//
// 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
//
// https://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.
#include <cstddef>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <memory>
#include <string>
#include <string_view>
#include <vector>
#include <fstream>
#include <implot.h>
#include <mujoco/mujoco.h>
#include "third_party/mujoco/src/experimental/platform/hal/graphics_mode.h"
#include "third_party/mujoco/src/experimental/platform/hal/renderer.h"
#include "third_party/mujoco/src/experimental/platform/hal/window.h"
#include "structs.h"
#include <pybind11/eval.h>
#include <pybind11/pybind11.h>
#include <pybind11/pytypes.h>
#include <pybind11/stl.h>
static bool IsCuda() {
#ifdef CUDA
return true;
#else
return false;
#endif
}
static bool IsCrd() {
const char* display = getenv("DISPLAY");
return display ? strcmp(display, ":20") == 0 : false;
}
static std::vector<std::byte> LoadAsset(std::string_view path) {
std::string file_path = "assets/" +
std::string(path.substr(path.find(':') + 1));
std::ifstream file(file_path, std::ios::binary | std::ios::ate);
if (!file.is_open()) {
return {};
}
auto file_size = file.tellg();
file.seekg(0, std::ios::beg);
std::vector<std::byte> buffer(file_size);
if (!file.read(reinterpret_cast<char*>(buffer.data()), file_size)) {
return {};
}
return buffer;
}
// Holds loaded resource data for the MuJoCo resource provider.
struct ResourceData {
std::vector<std::byte> bytes;
};
class Viewer {
public:
Viewer(const std::string& title, int width, int height,
std::string graphics_mode_str) {
// Register resource providers for font and filament assets.
mjpResourceProvider resource_provider;
mjp_defaultResourceProvider(&resource_provider);
resource_provider.open = [](mjResource* resource) {
auto* data = new ResourceData();
data->bytes = LoadAsset(resource->name);
resource->data = data;
return static_cast<int>(data->bytes.size());
};
resource_provider.read = [](mjResource* resource, const void** buffer) {
auto* data = static_cast<ResourceData*>(resource->data);
*buffer = data->bytes.data();
return static_cast<int>(data->bytes.size());
};
resource_provider.close = [](mjResource* resource) {
delete static_cast<ResourceData*>(resource->data);
resource->data = nullptr;
};
resource_provider.prefix = "font";
mjp_registerResourceProvider(&resource_provider);
resource_provider.prefix = "filament";
mjp_registerResourceProvider(&resource_provider);
mujoco::platform::Window::Config config;
using GraphicsMode = mujoco::platform::GraphicsMode;
config.gfx_mode = mujoco::platform::GraphicsModeFromString(
graphics_mode_str, GraphicsMode::FilamentOpenGl);
window_ = std::make_unique<mujoco::platform::Window>("PyStudio " + title,
width, height, config);
ImPlot::CreateContext();
renderer_ = std::make_unique<mujoco::platform::Renderer>(
window_->GetNativeWindowHandle(), config.gfx_mode);
}
void InitRenderer(const mujoco::python::MjModelWrapper& model) {
renderer_->Init(model.get());
}
bool NewFrame() {
const mujoco::platform::Window::Status status = window_->NewFrame();
return status == mujoco::platform::Window::Status::kRunning;
}
intptr_t UploadImage(intptr_t tex_id, const std::string img, int width,
int height, int bpp) {
return renderer_->UploadImage(tex_id, (const std::byte*)img.data(), width,
height, bpp);
}
int RenderToTexture(const mujoco::python::MjModelWrapper& model,
mujoco::python::MjDataWrapper& data,
mujoco::python::MjvCameraWrapper& cam, int width,
int height, int tex_id) {
const int bytes_per_pixel = 3;
std::vector<std::byte> bytes(width * height * bytes_per_pixel);
renderer_->RenderToTexture(model.get(), data.get(), cam.get(), width,
height, bytes.data());
return renderer_->UploadImage(tex_id, bytes.data(), width, height,
bytes_per_pixel);
}
std::string GetDropFile() {
return window_->GetDropFile();
}
void Present(const mujoco::python::MjModelWrapper& model,
mujoco::python::MjDataWrapper& data,
mujoco::python::MjvPerturbWrapper& perturb,
mujoco::python::MjvCameraWrapper& camera,
mujoco::python::MjvOptionWrapper& vis_options,
const std::vector<uint8_t>& render_flags) {
const float width = window_->GetWidth();
const float height = window_->GetHeight();
const float scale = window_->GetScale();
if (mujoco::platform::IsHeadless(window_->GetGraphicsMode())) {
pixels_.resize(width * height * 3);
} else {
pixels_.clear();
}
// Update render flags before rendering.
mjtByte* flags = renderer_->GetRenderFlags();
for (size_t i = 0; i < mjNRNDFLAG && i < render_flags.size(); ++i) {
flags[i] = render_flags[i];
}
renderer_->Render(model.get(), data.get(), perturb.get(), camera.get(),
vis_options.get(), width * scale, height * scale,
pixels_);
window_->EndFrame();
window_->Present(pixels_);
}
private:
std::unique_ptr<mujoco::platform::Window> window_;
std::unique_ptr<mujoco::platform::Renderer> renderer_;
std::vector<std::byte> pixels_;
};
PYBIND11_MODULE(native_viewer_cc, m) {
pybind11::class_<Viewer>(m, "Viewer")
.def(pybind11::init<const std::string&, int, int, const std::string&>())
.def("InitRenderer", &Viewer::InitRenderer)
.def("NewFrame", &Viewer::NewFrame)
.def("Present", &Viewer::Present)
.def("UploadImage", &Viewer::UploadImage)
.def("RenderToTexture", &Viewer::RenderToTexture)
.def("GetDropFile", &Viewer::GetDropFile);
m.def("IsCrd", &IsCrd);
m.def("IsCuda", &IsCuda);
}
python/mujoco/experimental/studio/native_viewer.py
+171
View File
@@ -0,0 +1,171 @@
# Copyright 2026 DeepMind Technologies Limited
#
# 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
#
# https://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.
"""Simulation-agnostic native viewer for MuJoCo models.
This class is simulation-agnostic and as such it does not own the model or data.
See the documentation for studio_app.py for more details on the architecture
separating the viewer and simulation. See the sample/ folder for examples of how
to use these classes.
"""
import mujoco
from mujoco.experimental.studio import native_viewer_cc as _viewer
from mujoco.experimental.studio import ux
class NativeViewer:
"""Simulation-agnostic native viewer for MuJoCo models."""
def __init__(
self,
model: mujoco.MjModel,
camera: mujoco.MjvCamera | None = None,
vis_options: mujoco.MjvOption | None = None,
perturb: mujoco.MjvPerturb | None = None,
render_flags: ux.RenderFlags | None = None,
title: str = '',
width: int = 1200,
height: int = 800,
gfx: str = '',
) -> None:
"""Initializes the NativeViewer.
The viewer creates and modifies its own camera, perturbation, and
visualization option objects unless they are provided.
Args:
model: The MuJoCo model, used to initialize the renderer.
camera: Camera parameters. Internal object is created if None.
vis_options: Visualization options. Internal object is created if None.
perturb: Perturbation parameters. Internal object is created if None.
render_flags: Render flags. Internal object is created if None.
title: Title of the viewer window.
width: Initial width of the viewer window.
height: Initial height of the viewer window.
gfx: Graphics mode.
"""
self.camera = camera or mujoco.MjvCamera()
self.perturb = perturb or mujoco.MjvPerturb()
self.vis_options = vis_options or mujoco.MjvOption()
self._viewer = _viewer.Viewer(title, width, height, gfx)
self._viewer.InitRenderer(model)
# This class does not own the model but we need to know if the model being
# rendered has changed, so we store the unique python object id here so we
# can use it to detect model changes.
self._renderer_model_id = id(model)
self._is_running = True
if render_flags is not None:
self.render_flags = render_flags
else:
self.render_flags = ux.RenderFlags()
# Initted to match mujoco/src/engine/engine_vis_init.c
self.render_flags.flags = [1, 0, 1, 0, 1, 0, 1, 0, 0, 0, 1]
def _sync_renderer(self, model: mujoco.MjModel) -> None:
"""Re-initializes the renderer if the model object has changed."""
if id(model) != self._renderer_model_id:
self._viewer.InitRenderer(model)
self._renderer_model_id = id(model)
def is_running(self) -> bool:
"""Poll for a new frame; returns ``False`` when the window is closed."""
if not self._is_running:
return False
self._is_running = self._viewer.NewFrame()
return self._is_running
def sync(
self,
model: mujoco.MjModel,
data: mujoco.MjData,
) -> None:
"""Render the scene and present it to the window.
Args:
model: The MuJoCo model provided by the simulation.
data: The MuJoCo data provided by the simulation.
"""
self._sync_renderer(model)
self._viewer.Present(
model,
data,
self.perturb,
self.camera,
self.vis_options,
self.render_flags.flags,
)
def stop(self) -> None:
"""Stop the viewer."""
self._is_running = False
def get_drop_file(self) -> str:
"""Returns the path of the file dropped into the window, or empty string."""
return self._viewer.GetDropFile()
def upload_image(
self, tex_id: int, img: str | bytes, width: int, height: int, bpp: int
) -> int:
"""Uploads an image to the backend for GUI rendering.
The ID can be used in subsequent calls to update the texture data. An empty
`img` argument will free the texture if it exists. A `tex_id` of 0 will
create a new texture.
Args:
tex_id: The texture ID.
img: The image data as string or bytes.
width: Width of the image.
height: Height of the image.
bpp: Bytes per pixel.
Returns:
The texture ID.
"""
return self._viewer.UploadImage(tex_id, img, width, height, bpp)
def render_to_texture(
self,
model: mujoco.MjModel,
data: mujoco.MjData,
tex_id: int,
width: int,
height: int,
) -> int:
"""Renders the scene to a texture.
This function renders the scene from the current camera view into a texture.
It handles buffer allocation internally.
Args:
model: The MuJoCo model provided by the simulation.
data: The MuJoCo data provided by the simulation.
tex_id: The texture ID to render into (0 to create a new one).
width: Width of the texture.
height: Height of the texture.
Returns:
The texture ID.
"""
self._sync_renderer(model)
return self._viewer.RenderToTexture(
model,
data,
self.camera,
width,
height,
tex_id,
)
python/mujoco/experimental/studio/parser.cc
+46
View File
@@ -0,0 +1,46 @@
// Copyright 2026 DeepMind Technologies Limited
//
// 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
//
// https://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.
#include <string_view>
#include <mujoco/mujoco.h>
#include "third_party/mujoco/src/experimental/platform/sim/model_holder.h"
#include "structs.h"
#include <pybind11/pybind11.h>
namespace mujoco::python {
// Loads, parses, and compiles a MuJoCo model from the given file. Returns the
// python mjData object (which also contains the compiled mjModel).
py::object Parse(std::string_view filepath) {
auto holder = platform::ModelHolder::FromFile(filepath);
if (!holder->ok()) {
throw py::value_error(
std::string("Failed to load model from '") +
std::string(filepath) + "': " + std::string(holder->error()));
}
mjModel* model = holder->ReleaseModel();
mjData* data = holder->ReleaseData();
py::object py_model = py::cast(MjModelWrapper(model));
py::object py_data =
py::cast(MjDataWrapper(py::cast<MjModelWrapper*>(py_model), data));
return py_data;
}
} // namespace mujoco::python
PYBIND11_MODULE(parser, m) {
m.def("parse", &mujoco::python::Parse,
pybind11::return_value_policy::take_ownership);
}
python/mujoco/experimental/studio/renderer.cc
+78
View File
@@ -0,0 +1,78 @@
// Copyright 2026 DeepMind Technologies Limited
//
// 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
//
// https://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.
#include "third_party/mujoco/src/experimental/platform/hal/renderer.h"
#include <cstddef>
#include <memory>
#include <optional>
#include <string>
#include <vector>
#include <mujoco/mujoco.h>
#include "third_party/mujoco/src/experimental/platform/hal/graphics_mode.h"
#include "structs.h"
#include <pybind11/eval.h>
#include <pybind11/pybind11.h>
#include <pybind11/pytypes.h>
#include <pybind11/stl.h>
namespace mujoco::python {
class Renderer {
public:
using RendererImpl = mujoco::platform::Renderer;
using GraphicsMode = mujoco::platform::GraphicsMode;
Renderer(const std::string& graphics_mode_str) {
const GraphicsMode mode = mujoco::platform::GraphicsModeFromString(
graphics_mode_str, GraphicsMode::FilamentOpenGl);
impl_ = std::make_unique<RendererImpl>(nullptr, mode);
}
void Init(const MjModelWrapper& model) { impl_->Init(model.get()); }
pybind11::bytes Render(const MjModelWrapper& model, MjDataWrapper& data,
std::optional<MjvPerturbWrapper>& perturb,
std::optional<MjvCameraWrapper>& camera,
std::optional<MjvOptionWrapper>& vis_option, int width,
int height) {
std::vector<std::byte> pixels(width * height * 3);
impl_->Render(
model.get(), data.get(), perturb ? perturb.value().get() : nullptr,
camera ? camera.value().get() : nullptr,
vis_option ? vis_option.value().get() : nullptr, width, height, pixels);
return pybind11::bytes((const char*)pixels.data(), pixels.size());
}
pybind11::memoryview GetRenderFlags() {
return pybind11::memoryview::from_buffer(
impl_->GetRenderFlags(), {static_cast<pybind11::ssize_t>(mjNRNDFLAG)},
{sizeof(mjtByte)});
}
private:
std::unique_ptr<RendererImpl> impl_;
};
} // namespace mujoco::python
PYBIND11_MODULE(renderer, m) {
pybind11::class_<mujoco::python::Renderer>(m, "Renderer")
.def(pybind11::init<const std::string&>())
.def("Init", &mujoco::python::Renderer::Init)
.def("Render", &mujoco::python::Renderer::Render)
.def("get_render_flags", &mujoco::python::Renderer::GetRenderFlags,
pybind11::keep_alive<0, 1>());
}
python/mujoco/experimental/studio/sample/async.py
+231
View File
@@ -0,0 +1,231 @@
# Copyright 2026 DeepMind Technologies Limited
#
# 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
#
# https://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.
"""This script runs a simulation and viewer in separate processes communicating asynchronously.
In this example, we will run the viewer in an independent process communicating
via multiprocessing queues. Controls are provided to simulate network transit
latency and adjust the communication rates.
You must provide a mjcf model file via the first command-line argument.
"""
import dataclasses
import multiprocessing
import os
import sys
import time
from absl import app as absl_app
from absl import flags as absl_flags
import mujoco
from mujoco.experimental.studio import native_viewer as _viewer
from mujoco.experimental.studio import sim as _sim
from mujoco.experimental.studio import studio_app
from mujoco.experimental.studio import ux
import numpy as np
from mujoco.experimental.dear_imgui import dear_imgui as imgui
_GFX = absl_flags.DEFINE_string('gfx', '', 'Rendering graphics mode.')
_WIDTH = absl_flags.DEFINE_integer('width', 1200, 'Width of the output image.')
_HEIGHT = absl_flags.DEFINE_integer('height', 800, 'Height of the output image')
@dataclasses.dataclass
class SimToView:
"""A message sent from the simulation process to the viewer process."""
model: mujoco.MjModel | None = None
data: mujoco.MjData | None = None
state: np.ndarray | None = None
state_sig: int = 0
send_time: float = 0.0
@dataclasses.dataclass
class ViewToSim:
"""A message sent from the viewer process to the simulation process."""
state: np.ndarray | None = None
state_sig: int = 0
reset: bool = False
send_rate: float = 60.0
class Network:
"""Simulated networking parameters."""
def __init__(self) -> None:
self.transit_buffer = []
self.send_rate = 60.0
self.network_delay = 0.2
def get_arrived(self, q: multiprocessing.Queue) -> SimToView | None:
now = time.time()
while not q.empty():
self.transit_buffer.append(q.get())
arrived = None
while (
self.transit_buffer
and now >= self.transit_buffer[0].send_time + self.network_delay
):
arrived = self.transit_buffer.pop(0)
return arrived
def view(
sim_to_view: multiprocessing.Queue,
view_to_sim: multiprocessing.Queue,
) -> None:
"""Entry-point for process that renders the simulation."""
# Block until the first message (containing the model) arrives.
msg = sim_to_view.get()
assert msg.model is not None, 'First message must contain the MuJoCo model.'
title = os.path.basename(sys.argv[0])
xfrc_sig = int(mujoco.mjtState.mjSTATE_XFRC_APPLIED)
xfrc_size = mujoco.mj_stateSize(msg.model, xfrc_sig)
xfrc_state = np.zeros(xfrc_size, np.float64)
app = studio_app.StudioApp(msg.model, msg.data)
network = Network()
viewer = _viewer.NativeViewer(
app.model,
title=title,
width=_WIDTH.value,
height=_HEIGHT.value,
gfx=_GFX.value,
)
while viewer.is_running() and app.is_running():
# Determine which messages have arrived through the simulated network.
arrived = network.get_arrived(sim_to_view)
# Update the camera and compute the perturbation.
app.handle_mouse_events(viewer.camera, viewer.vis_options, viewer.perturb)
# Sync state from the backend if a new payload actually arrived.
if arrived is not None and arrived.state is not None:
mujoco.mj_setState(app.model, app.data, arrived.state, arrived.state_sig)
mujoco.mj_forward(app.model, app.data)
# Always apply the perturbation forces from the viewer.
app.apply_perturb(viewer.perturb)
# Transmit user interaction when we get a new state
if arrived is not None:
mujoco.mj_getState(app.model, app.data, xfrc_state, xfrc_sig)
view_to_sim.put(
ViewToSim(
send_rate=network.send_rate, state=xfrc_state, state_sig=xfrc_sig
)
)
# Build the UI.
ux.setup_theme(app.theme)
if imgui.Begin(
'Settings',
flags=int(imgui.WindowFlags.AlwaysAutoResize)
| int(imgui.WindowFlags.NoTitleBar)
| int(imgui.WindowFlags.NoCollapse),
):
imgui.PushItemWidth(200.0)
_, network.network_delay = imgui.SliderFloat(
'Network Latency (s)', network.network_delay, 0.0, 2.0
)
updated, network.send_rate = imgui.SliderFloat(
'Send Rate (Hz)', network.send_rate, 1.0, 120.0
)
if updated:
view_to_sim.put(ViewToSim(send_rate=network.send_rate))
imgui.SetNextItemWidth(-1)
if imgui.Button('Reset Simulation'):
view_to_sim.put(ViewToSim(reset=True, send_rate=network.send_rate))
imgui.PopItemWidth()
imgui.End()
viewer.sync(app.model, app.data)
def sim(
data: mujoco.MjData,
model: mujoco.MjModel,
sim_to_view: multiprocessing.Queue,
view_to_sim: multiprocessing.Queue,
view_process: multiprocessing.Process,
) -> None:
"""Entry-point for process that runs the simulation."""
sim_to_view.put(SimToView(model=model, data=data))
step_control = _sim.StepControl()
integration_sig = int(mujoco.mjtState.mjSTATE_INTEGRATION)
integration_size = mujoco.mj_stateSize(model, integration_sig)
integration_state = np.empty(integration_size, np.float64)
msg = ViewToSim()
last_send_time = time.time()
while view_process.is_alive():
while not view_to_sim.empty():
msg = view_to_sim.get()
if msg.reset:
mujoco.mj_resetData(model, data)
mujoco.mj_forward(model, data)
msg.reset = False
# Apply perturbation forces received from the viewer process.
if msg.state is not None:
mujoco.mj_setState(model, data, msg.state, msg.state_sig)
# Advance the simulation keeping up with real-time.
step_control.advance(model, data)
# Send the simulation state paced by the requested send_rate.
now = time.time()
if now - last_send_time >= 1.0 / max(1.0, msg.send_rate):
mujoco.mj_getState(model, data, integration_state, integration_sig)
sim_to_view.put(
SimToView(
state=integration_state,
state_sig=integration_sig,
send_time=now,
)
)
last_send_time = now
def main(argv: list[str]) -> None:
app = studio_app.StudioApp.from_argv(argv)
# Queues for communication between the simulation and viewer processes.
sim_to_view = multiprocessing.Queue()
view_to_sim = multiprocessing.Queue()
# Start the viewer process.
view_process = multiprocessing.Process(
target=view, args=(sim_to_view, view_to_sim)
)
view_process.start()
# Start the simulation in the main process.
sim(app.data, app.model, sim_to_view, view_to_sim, view_process)
if __name__ == '__main__':
absl_app.run(main)
python/mujoco/experimental/studio/sample/implot.py
+199
View File
@@ -0,0 +1,199 @@
# Copyright 2026 DeepMind Technologies Limited
#
# 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
#
# https://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.
"""Example to run studio in the native viewer with responsive ImPlot UI.
This script runs a Studio viewer in-process and adds an 'Inspect Body' window
using ImGui and ImPlot bindings to visualize selected body data. The example
demonstrates how responsive UI layout rules are easily implemented.
Provide an MJCF model file via the first command-line argument to launch.
"""
import math
import os
import sys
from absl import app as absl_app
from absl import flags as absl_flags
import mujoco
from mujoco.experimental.studio import native_viewer as _viewer
from mujoco.experimental.studio import studio_app
import numpy as np
from mujoco.experimental.dear_imgui import dear_imgui as imgui
from mujoco.experimental.implot import implot
_GFX = absl_flags.DEFINE_string('gfx', '', 'Rendering graphics mode.')
_WIDTH = absl_flags.DEFINE_integer('width', 1200, 'Width of the output image.')
_HEIGHT = absl_flags.DEFINE_integer('height', 800, 'Height of the output image')
_N_HISTORY = 100
_PLOT_FLAGS = (
implot.Flags.NoInputs.value # Disable pan/zoom mouse interaction.
| implot.Flags.NoMenus.value # Disable right-click context menu.
| implot.Flags.NoBoxSelect.value # Disable drag-to-select regions.
)
_AXIS_FLAGS = (
implot.AxisFlags.NoGridLines.value # Hide background grid lines.
| implot.AxisFlags.NoTickMarks.value # Hide small tick marks on the axis.
)
def _setup_plot_flags(plot_size: imgui.Vec2) -> int:
flags = _PLOT_FLAGS
if min(plot_size.x, plot_size.y) < 300:
flags |= implot.Flags.NoTitle.value
if min(plot_size.x, plot_size.y) < 200:
flags |= implot.Flags.NoLegend.value
return flags
def _setup_time_axis(plot_size: imgui.Vec2) -> None:
flags = _AXIS_FLAGS
if plot_size.x < 300:
flags |= implot.AxisFlags.NoTickLabels.value
implot.SetupAxis(implot.Axis.X1, '', flags)
implot.SetupAxisLimits(implot.Axis.X1, 0, _N_HISTORY)
def _setup_xpos_axis(centroid: list[np.ndarray], plot_size: imgui.Vec2) -> None:
flags = _AXIS_FLAGS
if plot_size.y < 300:
flags |= implot.AxisFlags.NoTickLabels.value
implot.SetupAxis(implot.Axis.Y1, '', flags)
min_y = min(c[1] for c in centroid)
max_y = max(c[1] for c in centroid)
margin = max((max_y - min_y) * 0.1, 0.05)
implot.SetupAxisLimits(
implot.Axis.Y1,
min_y - margin,
max_y + margin,
cond=implot.Cond.Always,
)
def _setup_angle_axis(plot_size: imgui.Vec2) -> None:
flags = _AXIS_FLAGS
if plot_size.y < 300:
flags |= implot.AxisFlags.NoTickLabels.value
implot.SetupAxis(implot.Axis.Y1, '', flags)
implot.SetupAxisLimits(implot.Axis.Y1, -185.0, 185.0)
implot.SetupAxisTicks(
implot.Axis.Y1,
[-180.0, -90.0, 0.0, 90.0, 180.0],
['-180', '-90', '0', '90', '180'],
)
def main(argv: list[str]) -> None:
app = studio_app.StudioApp.from_argv(argv)
title = os.path.basename(sys.argv[0])
# Initialize the viewer.
viewer = _viewer.NativeViewer(
app.model,
title=title,
width=_WIDTH.value,
height=_HEIGHT.value,
gfx=_GFX.value,
)
# Variables for the custom UI.
centroid = [np.zeros(3) for _ in range(_N_HISTORY)]
euler = [np.zeros(3) for _ in range(_N_HISTORY)]
body_id = -1
# Main viewer loop.
while viewer.is_running():
if not app.update(viewer.camera, viewer.vis_options, viewer.perturb):
break
# Build standard Studio UI.
app.build_gui(viewer.camera, viewer.vis_options, viewer.render_flags)
# Inspect the perturb.select body
if viewer.perturb.select > 0:
body_id = viewer.perturb.select
# Display selected body information.
if body_id > 0:
body_name = mujoco.mj_id2name(
app.model, int(mujoco.mjtObj.mjOBJ_BODY), body_id
)
imgui.SetNextWindowSize(imgui.Vec2(1200, 600), imgui.Cond.FirstUseEver)
# Note: The window title uses the special "###" markup to ensure the imgui
# ID for the window is constant for all body names. This is needed for
# the window to retain its state for all bodies.
window_title = f'Inspect Body {body_name or "(???)"!r} ({body_id})###Plot'
if imgui.Begin(window_title):
avail = imgui.GetContentRegionAvail()
wide = avail.x > avail.y
# Add a small padding factor to prevent scrollbars.
plot_size = imgui.Vec2(
avail.x * 0.5 - 4 if wide else avail.x,
avail.y if wide else avail.y * 0.5 - 4,
)
plot_flags = _setup_plot_flags(plot_size)
if implot.BeginPlot('Centroid vs Time', plot_size, flags=plot_flags):
_setup_time_axis(plot_size)
_setup_xpos_axis(centroid, plot_size)
implot.PlotLine('x', range(_N_HISTORY), [c[0] for c in centroid])
implot.PlotLine('y', range(_N_HISTORY), [c[1] for c in centroid])
implot.PlotLine('z', range(_N_HISTORY), [c[2] for c in centroid])
implot.EndPlot()
if wide:
imgui.SameLine()
if implot.BeginPlot('Euler Angle vs Time', plot_size, flags=plot_flags):
_setup_time_axis(plot_size)
_setup_angle_axis(plot_size)
implot.PlotLine('roll', range(_N_HISTORY), [e[0] for e in euler])
implot.PlotLine('pitch', range(_N_HISTORY), [e[1] for e in euler])
implot.PlotLine('yaw', range(_N_HISTORY), [e[2] for e in euler])
implot.EndPlot()
imgui.End()
# Update plot data
centroid.pop(0)
euler.pop(0)
if body_id > 0:
centroid.append(app.data.xpos[body_id].copy())
# Convert quaternion to Euler angles via rotation matrix.
quat = app.data.xquat[body_id]
mat = np.zeros(9)
mujoco.mju_quat2Mat(mat, quat)
# mat is row-major 3x3: R[i,j] = mat[3*i + j].
roll = math.atan2(mat[7], mat[8])
pitch = math.atan2(-mat[6], math.sqrt(mat[7] ** 2 + mat[8] ** 2))
yaw = math.atan2(mat[3], mat[0])
euler.append(np.degrees(np.array([roll, pitch, yaw])))
else:
centroid.append(np.zeros(3))
euler.append(np.zeros(3))
viewer.sync(app.model, app.data)
viewer.stop()
if __name__ == '__main__':
absl_app.run(main)
python/mujoco/experimental/studio/sample/render.py
+73
View File
@@ -0,0 +1,73 @@
# Copyright 2026 DeepMind Technologies Limited
#
# 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
#
# https://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.
"""Render a MuJoCo model to an image."""
import os
import sys
from absl import app
from absl import flags
import mujoco
from mujoco.experimental.studio import parser
from mujoco.experimental.studio import renderer
from PIL import Image
_MODEL = flags.DEFINE_string('model', '', 'Model file to load.')
_OUTPUT = flags.DEFINE_string('output', '', 'Output file to save.')
_GFX = flags.DEFINE_string('gfx', '', 'Renderer to use.')
_WIDTH = flags.DEFINE_integer('width', 320, 'Width of the output image.')
_HEIGHT = flags.DEFINE_integer('height', 240, 'Height of the output image.')
_STEPS = flags.DEFINE_integer('steps', 1, 'Number of steps before render.')
def main(argv):
if len(argv) > 1:
raise app.UsageError('Too many command-line arguments.')
if not _MODEL.value:
raise ValueError('`model` flag is required.')
if not _OUTPUT.value:
raise ValueError('`output flag is required.')
try:
data = parser.parse(_MODEL.value)
model = data.model
except Exception as ex: # pylint: disable=broad-except
print(f'Error loading model from `{_MODEL.value}`: {ex}')
sys.exit(-1)
for _ in range(_STEPS.value):
mujoco.mj_step(model, data)
try:
r = renderer.Renderer(_GFX.value)
r.Init(model)
pixels = r.Render(
model, data, None, None, None, _WIDTH.value, _HEIGHT.value
)
except Exception as ex: # pylint: disable=broad-except
print(f'Error rendering model: {ex}')
sys.exit(-2)
try:
img = Image.frombytes('RGB', (_WIDTH.value, _HEIGHT.value), pixels)
img.save(_OUTPUT.value, format=os.path.splitext(_OUTPUT.value)[1][1:])
except Exception as ex: # pylint: disable=broad-except
print(f'Error saving image to `{_OUTPUT.value}`: {ex}')
sys.exit(-3)
return 0
if __name__ == '__main__':
app.run(main)
python/mujoco/experimental/studio/sim.cc
+75
View File
@@ -0,0 +1,75 @@
// Copyright 2026 DeepMind Technologies Limited
//
// 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
//
// https://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.
// Python bindings for MuJoCo platform simulation components.
#include <mujoco/mujoco.h>
#include "third_party/mujoco/src/experimental/platform/sim/step_control.h"
#include "structs.h"
#include <pybind11/pybind11.h>
namespace py = pybind11;
using StepControl = mujoco::platform::StepControl;
PYBIND11_MODULE(sim, m) {
m.doc() = "MuJoCo platform simulation bindings for Link.";
py::enum_<StepControl::Status>(m, "StepStatus")
.value("OK", StepControl::Status::kOk)
.value("PAUSED", StepControl::Status::kPaused)
.value("VISCOUS_PAUSED", StepControl::Status::kViscousPaused)
.value("AUTO_RESET", StepControl::Status::kAutoReset)
.value("DIVERGED", StepControl::Status::kDiverged);
py::enum_<StepControl::PauseState>(m, "PauseState")
.value("UNPAUSED", StepControl::PauseState::kUnpaused)
.value("NORMAL_PAUSED", StepControl::PauseState::kNormalPaused)
.value("VISCOUS_PAUSED", StepControl::PauseState::kViscousPaused);
py::class_<StepControl>(m, "StepControl")
.def(py::init<>())
.def(
"advance",
[](StepControl& self, py::object model_obj, py::object data_obj,
py::object step_fn) {
auto& model = py::cast<mujoco::python::MjModelWrapper&>(model_obj);
auto& data = py::cast<mujoco::python::MjDataWrapper&>(data_obj);
if (step_fn.is_none()) {
return self.Advance(model.get(), data.get());
} else {
return self.Advance(
model.get(), data.get(),
[step_fn, model_obj, data_obj](mjModel*, mjData*) {
step_fn(model_obj, data_obj);
});
}
},
py::arg("model"), py::arg("data"), py::arg("step_fn") = py::none(),
"Step physics forward, respecting speed settings and refresh budget.")
.def("force_sync", &StepControl::ForceSync,
"Ensures the next Advance() will synchronize time and step once.")
.def("get_speed", &StepControl::GetSpeed,
"Returns the desired simulation speed as a percentage of real time.")
.def("get_speed_measured", &StepControl::GetSpeedMeasured,
"Returns the measured simulation speed.")
.def("set_speed", &StepControl::SetSpeed, py::arg("speed"),
"Sets the desired speed (clamped to [0.1%, 100%]).")
.def("set_pause_state", &StepControl::SetPauseState, py::arg("state"),
"Sets the pause state of the simulation.")
.def("get_pause_state", &StepControl::GetPauseState,
"Returns the current pause state.")
.def("request_single_step", &StepControl::RequestSingleStep,
"Request a single step if paused.");
}
python/mujoco/experimental/studio/studio.py
+50
View File
@@ -0,0 +1,50 @@
# Copyright 2026 DeepMind Technologies Limited
#
# 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
#
# https://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.
"""This script runs Studio from Python, visualized in a native viewer."""
from absl import app as absl_app
from absl import flags as absl_flags
from mujoco.experimental.studio import native_viewer
from mujoco.experimental.studio import studio_app
_GFX = absl_flags.DEFINE_string('gfx', '', 'Rendering graphics mode.')
_WIDTH = absl_flags.DEFINE_integer('width', 1200, 'Width of the output image.')
_HEIGHT = absl_flags.DEFINE_integer('height', 800, 'Height of the output image')
def main(argv: list[str]) -> None:
app = studio_app.StudioApp.from_argv(argv)
# Initialize the viewer.
viewer = native_viewer.NativeViewer(
app.model,
width=_WIDTH.value,
height=_HEIGHT.value,
gfx=_GFX.value,
)
# Main viewer loop.
while viewer.is_running():
if not app.update_from_viewer(viewer):
break
app.build_gui(viewer.camera, viewer.vis_options, viewer.render_flags)
viewer.sync(app.model, app.data)
viewer.stop()
if __name__ == '__main__':
absl_app.run(main)
python/mujoco/experimental/studio/studio_app.py
+478
View File
@@ -0,0 +1,478 @@
# Copyright 2026 DeepMind Technologies Limited
#
# 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
#
# https://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.
"""Viewer-agnostic Python implementation of Studio.
Architecture:
StudioApp owns the simulation state (model, data) and the UI logic.
Viewer classes (e.g., NativeViewer) own the window (if required), renderer,
camera, and visualization options. The viewer never stores references to
model or data. Instead, the caller passes them each frame via
viewer.sync(model, data). This ensures the viewer always renders the current
model, even if StudioApp.load_model_from_file() swaps it.
The class can be used to implement the full Studio application in Python. By
using the more granular member functions it can also build simple apps that only
use a subset of the Studio UI. This configuration is fully dynamic, there is
nothing to configure in advance, you can change your app by changing the
functions that get called each frame. This class is also viewer-agnostic and as
such does not own camera, vis_options or perturb objects (these are provided by
the viewer).
See the sample/ folder for usage examples.
"""
import os
import sys
import typing
import mujoco
from mujoco.experimental.studio import parser
from mujoco.experimental.studio import sim
from mujoco.experimental.studio import studio_app_events as events
from mujoco.experimental.studio import ux
from mujoco.experimental.studio import viewer_protocol
import numpy as np
from mujoco.experimental.dear_imgui import dear_imgui as imgui
# Type alias for a custom physics step function.
StepFn = typing.Callable[[mujoco.MjModel, mujoco.MjData], None]
def load_model_from_file(
model_path: str,
) -> tuple[mujoco.MjModel, mujoco.MjData] | None:
"""Loads a model and data from a file path."""
try:
data = parser.parse(model_path)
return data.model, data
except Exception as ex: # pylint: disable=broad-except
print(f'Error loading model from {model_path!r}: {ex}')
return None
class StudioApp:
"""Viewer-agnostic Python implementation of Studio."""
@classmethod
def from_argv(cls, argv: list[str]) -> 'StudioApp':
"""Constructs a StudioApp by parsing a model path from command-line args."""
if len(argv) < 2:
model = mujoco.MjSpec().compile()
data = mujoco.MjData(model)
app = cls(model, data)
app.step_control.set_pause_state(sim.PauseState.NORMAL_PAUSED)
return app
model_path = argv[1]
res = load_model_from_file(model_path)
if res is None:
sys.exit(-1)
model, data = res
app = cls(model, data)
app.model_path = model_path
return app
def load_model_from_file(
self, model_path: str
) -> tuple[mujoco.MjModel, mujoco.MjData] | None:
"""Loads a new model from a file, replacing the current model and data."""
res = load_model_from_file(model_path)
if res is None:
self.status = f'Error loading model from {model_path!r}'
return None
model, data = res
self.model = model
self.data = data
self.model_path = model_path
self.step_control = sim.StepControl()
self.ux_state = ux.UxState()
self.status = f'Loaded: {os.path.basename(model_path)!r}'
return model, data
def __init__(
self,
model: mujoco.MjModel,
data: mujoco.MjData,
):
"""Initializes the Studio application."""
self.model = model
self.data = data
self.model_path = ''
self.step_control = sim.StepControl()
self.ux_state = ux.UxState()
self.theme = ux.GuiTheme.LIGHT
self.show_stats = False
self.show_solver = False
self.should_quit = False
self.status = 'Ready'
# TODO(matijak): This should be part of the viewer, also making a struct to
# pass it around with the camera would be convenient.
self._cam_speed = 0.001
def handle_vis_options_keyboard_events(
self,
vis_options: mujoco.MjvOption,
is_freecam_wasd: bool,
) -> bool:
"""Toggles visualization flags based on keyboard shortcuts.
Args:
vis_options: The visualization options to modify.
is_freecam_wasd: If True, keys Q/E/A/D are reserved for camera movement
and will not toggle visualization flags.
Returns:
True if a key was handled, False otherwise.
"""
if imgui.GetIO().WantCaptureKeyboard:
return False
return events.handle_vis_options_keyboard_events(
vis_options, is_freecam_wasd
)
def handle_step_control_keyboard_events(self) -> bool:
"""Handles keyboard shortcuts for simulation stepping control.
Returns:
True if a key was handled, False otherwise.
"""
if imgui.GetIO().WantCaptureKeyboard:
return False
return events.handle_step_control_keyboard_events(
self.model, self.data, self.step_control, self.ux_state
)
def handle_freecam_wasd_keyboard_events(
self,
camera: mujoco.MjvCamera,
) -> bool:
"""Handles keyboard shortcuts for free camera movement."""
if imgui.GetIO().WantCaptureKeyboard:
return False
handled, self._cam_speed = events.handle_freecam_wasd_keyboard_events(
self.model, self.data, camera, self._cam_speed
)
return handled
def handle_keyboard_events(
self,
camera: mujoco.MjvCamera,
vis_options: mujoco.MjvOption,
) -> bool:
"""Handle keyboard events according to Studio's bindings."""
if imgui.GetIO().WantCaptureKeyboard:
return False
is_freecam_wasd = self.ux_state.camera_index == ux.FREE_CAMERA_IDX
if events.handle_step_control_keyboard_events(
self.model, self.data, self.step_control, self.ux_state
):
return True
if events.handle_camera_select_keyboard_events(
self.model, camera, self.ux_state
):
return True
if events.handle_vis_options_keyboard_events(vis_options, is_freecam_wasd):
return True
if is_freecam_wasd:
handled, self._cam_speed = events.handle_freecam_wasd_keyboard_events(
self.model, self.data, camera, self._cam_speed
)
if handled:
return True
return False
def handle_camera_tracking_mouse_events(
self,
camera: mujoco.MjvCamera,
vis_options: mujoco.MjvOption,
) -> None:
"""Handles mouse events for camera tracking."""
if imgui.GetIO().WantCaptureMouse:
return
events.handle_camera_tracking_mouse_events(
self.model, self.data, camera, vis_options, self.ux_state
)
def handle_mouse_events(
self,
camera: mujoco.MjvCamera,
vis_options: mujoco.MjvOption,
perturb: mujoco.MjvPerturb,
) -> None:
"""Handles mouse events."""
if imgui.GetIO().WantCaptureMouse:
return
events.handle_mouse_events(
self.model, self.data, camera, vis_options, perturb, self.ux_state
)
def reset_physics(self) -> None:
"""Reset the physics."""
mujoco.mj_resetData(self.model, self.data)
mujoco.mj_forward(self.model, self.data)
def apply_perturb(self, perturb: mujoco.MjvPerturb) -> None:
"""Apply perturbation the model."""
if self.step_control.get_pause_state() != sim.PauseState.NORMAL_PAUSED:
sig = mujoco.mjtState.mjSTATE_XFRC_APPLIED.value
size = mujoco.mj_stateSize(self.model, sig)
zero_state = np.zeros(size, np.float64)
mujoco.mj_setState(self.model, self.data, zero_state, sig)
mujoco.mjv_applyPerturbPose(self.model, self.data, perturb, 0)
mujoco.mjv_applyPerturbForce(self.model, self.data, perturb)
else:
mujoco.mjv_applyPerturbPose(self.model, self.data, perturb, 1)
def update_physics(
self,
perturb: mujoco.MjvPerturb,
*,
step_fn: StepFn | None = None,
) -> None:
"""Applies the perturbations and advances the physics.
Args:
perturb: The MuJoCo perturbation object.
step_fn: Optional custom physics step function. When provided, it is
called instead of ``step_control.advance``. The function receives
``(model, data)`` and should step the simulation in-place.
"""
self.apply_perturb(perturb)
if step_fn is not None:
step_fn(self.model, self.data)
else:
advance_status = self.step_control.advance(self.model, self.data)
if advance_status == sim.StepStatus.AUTO_RESET:
self.reset_physics()
def reset_physics_gui(self) -> None:
"""GUI to Reset the physics i.e., the reset button."""
button_size = imgui.GetFrameHeight()
square_size = imgui.Vec2(button_size, button_size)
icon_reset_model = '\uf0e2' # FontAwesome "undo" icon.
if imgui.Button(icon_reset_model, square_size):
self.reset_physics()
imgui.SetItemTooltip('Reset')
def is_running(self) -> bool:
"""Returns True if the application should continue running (called by update())."""
return not self.should_quit
def update(
self,
camera: mujoco.MjvCamera,
vis_options: mujoco.MjvOption,
perturb: mujoco.MjvPerturb,
*,
drop_file: str = '',
step_fn: StepFn | None = None,
) -> bool:
"""Update the simulation and handle user input.
Handles mouse input to compute perturbations or camera motion.
Handles keyboard input e.g., for keybindings or camera motion.
Applies the perturbations and advances the physics.
The argument objects are provided by the viewer.
Args:
camera: The MuJoCo camera object.
vis_options: The MuJoCo visualization options.
perturb: The MuJoCo perturbation object.
drop_file: Path of a file dropped into the viewer window. If non-empty the
current model is replaced with the dropped file.
step_fn: Optional custom physics step function. When provided, it is
called instead of ``step_control.advance``.
Returns:
Whether the application should continue running, this is a
convenience to allow this function to be used in a while loop.
"""
if drop_file:
self.load_model_from_file(drop_file)
self.handle_mouse_events(camera, vis_options, perturb)
self.handle_keyboard_events(camera, vis_options)
self.update_physics(perturb, step_fn=step_fn)
return self.is_running()
def update_from_viewer(
self,
viewer: viewer_protocol.Viewer,
*,
step_fn: StepFn | None = None,
) -> bool:
"""Convenience wrapper around update() that unpacks viewer attributes.
Args:
viewer: A viewer conforming to the Viewer protocol.
step_fn: Optional custom physics step function. When provided, it is
called instead of ``step_control.advance``.
Returns:
Whether the application should continue running.
"""
return self.update(
viewer.camera,
viewer.vis_options,
viewer.perturb,
drop_file=viewer.get_drop_file(),
step_fn=step_fn,
)
def build_gui(
self,
camera: mujoco.MjvCamera,
vis_options: mujoco.MjvOption,
render_flags: ux.RenderFlags,
) -> None:
"""Emit full Studio UI."""
ux.setup_theme(self.theme)
ux.configure_docking_layout()
# -- Main menu bar --------------------------------------------------------
if imgui.BeginMainMenuBar():
if imgui.BeginMenu('File'):
if imgui.MenuItem('Quit'):
self.should_quit = True
imgui.EndMenu()
if imgui.BeginMenu('Simulation'):
imgui.EndMenu()
if imgui.BeginMenu('Charts'):
if imgui.MenuItem('Solver', '', self.show_solver):
self.show_solver = not self.show_solver
if imgui.MenuItem('Stats', '', self.show_stats):
self.show_stats = not self.show_stats
imgui.EndMenu()
if imgui.BeginMenu('Help'):
if imgui.MenuItem('Stats', '', self.show_stats):
self.show_stats = not self.show_stats
imgui.Separator()
version = f'Version {mujoco.mj_versionString()}'
imgui.MenuItem(version)
imgui.EndMenu()
imgui.EndMainMenuBar()
# -- Tool Bar -------------------------------------------------------------
if imgui.Begin('ToolBar'):
imgui.PushStyleVar(imgui.StyleVar.CellPadding, imgui.Vec2(0, 0))
if imgui.BeginTable('##ToolBarTable', 2):
imgui.TableSetupColumn('', int(imgui.TableColumnFlags.WidthStretch))
imgui.TableSetupColumn('', int(imgui.TableColumnFlags.WidthFixed))
imgui.TableNextColumn()
self.reset_physics_gui()
imgui.SameLine()
ux.step_control_gui(self.model, self.step_control, self.ux_state)
imgui.TableNextColumn()
ux.camera_selection_gui(self.model, self.data, camera, self.ux_state)
imgui.SameLine()
ux.label_selection_gui(vis_options)
imgui.SameLine()
ux.frame_selection_gui(vis_options)
imgui.SameLine()
changed, self.theme = ux.theme_select_gui(self.theme)
if changed:
ux.setup_theme(self.theme)
imgui.EndTable()
imgui.PopStyleVar()
imgui.End()
# -- Left pane: Options ---------------------------------------------------
node_flags = int(imgui.TreeNodeFlags.SpanAvailWidth) | int(
imgui.TreeNodeFlags.Framed
)
imgui.Begin('Options')
if imgui.TreeNodeEx('Physics Settings', node_flags):
ux.physics_gui(self.model)
imgui.TreePop()
if imgui.TreeNodeEx('Rendering Settings', node_flags):
ux.rendering_gui(self.model, vis_options, render_flags)
imgui.TreePop()
if imgui.TreeNodeEx('Visibility Groups', node_flags):
ux.groups_gui(self.model, vis_options)
imgui.TreePop()
if imgui.TreeNodeEx('Visualization', node_flags):
ux.visualization_gui(self.model, vis_options, camera)
imgui.TreePop()
imgui.End()
# -- Right pane: Inspector ------------------------------------------------
imgui.Begin('Inspector')
if imgui.TreeNodeEx('Noise', node_flags):
ux.noise_gui(self.model, self.data, self.ux_state)
imgui.TreePop()
if imgui.TreeNodeEx('Joints', node_flags):
ux.joints_gui(self.model, self.data, vis_options)
imgui.TreePop()
if imgui.TreeNodeEx('Controls', node_flags):
ux.controls_gui(self.model, self.data, vis_options)
imgui.TreePop()
if imgui.TreeNodeEx(
'Sensors', node_flags | int(imgui.TreeNodeFlags.DefaultOpen)
):
ux.sensor_gui(self.model, self.data)
imgui.TreePop()
if imgui.TreeNodeEx('Watch', node_flags):
ux.watch_gui(self.model, self.data, self.ux_state)
imgui.TreePop()
if imgui.TreeNodeEx('State', node_flags):
ux.state_gui(self.model, self.data, self.ux_state)
imgui.TreePop()
imgui.End()
# -- Floating windows -----------------------------------------------------
if self.show_solver:
_, self.show_solver = imgui.Begin('Solver', self.show_solver)
ux.counts_gui(self.model, self.data)
ux.convergence_gui(self.model, self.data)
imgui.End()
if self.show_stats:
_, self.show_stats = imgui.Begin('Stats', self.show_stats)
paused = self.step_control.get_pause_state() != sim.PauseState.UNPAUSED
ux.stats_gui(self.model, self.data, paused, 0.0)
imgui.End()
# -- Status bar -----------------------------------------------------------
imgui.PushStyleVar(imgui.StyleVar.CellPadding, imgui.Vec2(0, 0))
imgui.PushStyleVar(imgui.StyleVar.FramePadding, imgui.Vec2(0, 0))
imgui.PushStyleVar(imgui.StyleVar.WindowPadding, imgui.Vec2(0, 0))
if imgui.Begin('StatusBar'):
imgui.Text(self.status)
imgui.End()
imgui.PopStyleVar(3)
python/mujoco/experimental/studio/studio_app_events.py
+592
View File
@@ -0,0 +1,592 @@
# Copyright 2026 DeepMind Technologies Limited
#
# 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
#
# https://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.
"""Temporary event handling functions for StudioApp."""
# TODO(matijak): These free functions implement the keyboard and mouse event
# handling for Studio. They are separated from the main StudioApp class to keep
# clarify the long-term API and avoid cluttering it with a large amount of
# temporary code. When studio/platform has a proper API for registering key
# bindings and mouse behaviour, the event handling functions will delegate to
# code shared with the C++ studio application.
import mujoco
from mujoco.experimental.studio import sim
from mujoco.experimental.studio import ux
import numpy as np
from mujoco.experimental.dear_imgui import dear_imgui as imgui
def handle_vis_options_keyboard_events(
vis_options: mujoco.MjvOption,
is_freecam_wasd: bool,
) -> bool:
"""Toggles visualization flags based on keyboard shortcuts.
Args:
vis_options: The visualization options to modify.
is_freecam_wasd: If True, keys Q/E/A/D are reserved for camera movement and
will not toggle visualization flags.
Returns:
True if a key was handled, False otherwise.
"""
if imgui.GetIO().WantCaptureKeyboard:
return False
pressed = imgui.IsKeyChordPressed
# Frame and label cycling.
if pressed(imgui.Key.F6):
vis_options.frame = (vis_options.frame + 1) % mujoco.mjtFrame.mjNFRAME.value
elif pressed(imgui.Key.F7):
vis_options.label = (vis_options.label + 1) % mujoco.mjtLabel.mjNLABEL.value
# Visualization flag toggles (single-key shortcuts).
elif pressed(imgui.Key.H):
vis_options.flags[mujoco.mjtVisFlag.mjVIS_CONVEXHULL] ^= 1
elif pressed(imgui.Key.X):
vis_options.flags[mujoco.mjtVisFlag.mjVIS_TEXTURE] ^= 1
elif pressed(imgui.Key.J):
vis_options.flags[mujoco.mjtVisFlag.mjVIS_JOINT] ^= 1
elif not is_freecam_wasd and pressed(imgui.Key.Q):
vis_options.flags[mujoco.mjtVisFlag.mjVIS_CAMERA] ^= 1
elif pressed(imgui.Key.U):
vis_options.flags[mujoco.mjtVisFlag.mjVIS_ACTUATOR] ^= 1
elif pressed(imgui.Key.Comma):
vis_options.flags[mujoco.mjtVisFlag.mjVIS_ACTIVATION] ^= 1
elif pressed(imgui.Key.Z):
vis_options.flags[mujoco.mjtVisFlag.mjVIS_LIGHT] ^= 1
elif pressed(imgui.Key.V):
vis_options.flags[mujoco.mjtVisFlag.mjVIS_TENDON] ^= 1
elif pressed(imgui.Key.Y):
vis_options.flags[mujoco.mjtVisFlag.mjVIS_RANGEFINDER] ^= 1
elif not is_freecam_wasd and pressed(imgui.Key.E):
vis_options.flags[mujoco.mjtVisFlag.mjVIS_CONSTRAINT] ^= 1
elif pressed(imgui.Key.I):
vis_options.flags[mujoco.mjtVisFlag.mjVIS_INERTIA] ^= 1
elif pressed(imgui.Key.Apostrophe):
vis_options.flags[mujoco.mjtVisFlag.mjVIS_SCLINERTIA] ^= 1
elif pressed(imgui.Key.B):
vis_options.flags[mujoco.mjtVisFlag.mjVIS_PERTFORCE] ^= 1
elif pressed(imgui.Key.O):
vis_options.flags[mujoco.mjtVisFlag.mjVIS_PERTOBJ] ^= 1
elif pressed(imgui.Key.C):
vis_options.flags[mujoco.mjtVisFlag.mjVIS_CONTACTPOINT] ^= 1
elif pressed(imgui.Key.N):
vis_options.flags[mujoco.mjtVisFlag.mjVIS_ISLAND] ^= 1
elif pressed(imgui.Key.F):
vis_options.flags[mujoco.mjtVisFlag.mjVIS_CONTACTFORCE] ^= 1
elif pressed(imgui.Key.P):
vis_options.flags[mujoco.mjtVisFlag.mjVIS_CONTACTSPLIT] ^= 1
elif pressed(imgui.Key.T):
vis_options.flags[mujoco.mjtVisFlag.mjVIS_TRANSPARENT] ^= 1
elif not is_freecam_wasd and pressed(imgui.Key.A):
vis_options.flags[mujoco.mjtVisFlag.mjVIS_AUTOCONNECT] ^= 1
elif pressed(imgui.Key.M):
vis_options.flags[mujoco.mjtVisFlag.mjVIS_COM] ^= 1
elif not is_freecam_wasd and pressed(imgui.Key.D):
vis_options.flags[mujoco.mjtVisFlag.mjVIS_STATIC] ^= 1
elif pressed(imgui.Key.Semicolon):
vis_options.flags[mujoco.mjtVisFlag.mjVIS_SKIN] ^= 1
elif pressed(imgui.Key.GraveAccent):
vis_options.flags[mujoco.mjtVisFlag.mjVIS_BODYBVH] ^= 1
elif pressed(imgui.Key.Backslash):
vis_options.flags[mujoco.mjtVisFlag.mjVIS_MESHBVH] ^= 1
# Site group toggles (Shift + 0-5).
elif pressed(int(imgui.Key.N0) | int(imgui.Key.Shift)):
vis_options.sitegroup[0] ^= 1
elif pressed(int(imgui.Key.N1) | int(imgui.Key.Shift)):
vis_options.sitegroup[1] ^= 1
elif pressed(int(imgui.Key.N2) | int(imgui.Key.Shift)):
vis_options.sitegroup[2] ^= 1
elif pressed(int(imgui.Key.N3) | int(imgui.Key.Shift)):
vis_options.sitegroup[3] ^= 1
elif pressed(int(imgui.Key.N4) | int(imgui.Key.Shift)):
vis_options.sitegroup[4] ^= 1
elif pressed(int(imgui.Key.N5) | int(imgui.Key.Shift)):
vis_options.sitegroup[5] ^= 1
# Geom group toggles (0-5).
elif pressed(imgui.Key.N0):
vis_options.geomgroup[0] ^= 1
elif pressed(imgui.Key.N1):
vis_options.geomgroup[1] ^= 1
elif pressed(imgui.Key.N2):
vis_options.geomgroup[2] ^= 1
elif pressed(imgui.Key.N3):
vis_options.geomgroup[3] ^= 1
elif pressed(imgui.Key.N4):
vis_options.geomgroup[4] ^= 1
elif pressed(imgui.Key.N5):
vis_options.geomgroup[5] ^= 1
else:
return False
return True
def handle_step_control_keyboard_events(
model: mujoco.MjModel,
data: mujoco.MjData,
step_control: sim.StepControl,
ux_state: ux.UxState,
) -> bool:
"""Handles keyboard shortcuts for simulation stepping control.
Args:
model: The MuJoCo model.
data: The MuJoCo data.
step_control: The simulation step control object.
ux_state: The UX state object.
Returns:
True if a key was handled, False otherwise.
"""
if imgui.GetIO().WantCaptureKeyboard:
return False
pressed = imgui.IsKeyChordPressed
if pressed(int(imgui.Key.Ctrl) | int(imgui.Key.Space)):
if step_control.get_pause_state() == sim.PauseState.VISCOUS_PAUSED:
step_control.set_pause_state(sim.PauseState.UNPAUSED)
else:
step_control.set_pause_state(sim.PauseState.VISCOUS_PAUSED)
return True
elif pressed(imgui.Key.Space):
pause = step_control.get_pause_state()
if pause in (sim.PauseState.VISCOUS_PAUSED, sim.PauseState.UNPAUSED):
step_control.set_pause_state(sim.PauseState.NORMAL_PAUSED)
else:
step_control.set_pause_state(sim.PauseState.UNPAUSED)
return True
elif pressed(imgui.Key.Backspace):
mujoco.mj_resetData(model, data)
mujoco.mj_forward(model, data)
return True
elif pressed(imgui.Key.Minus):
ux.set_speed_index(step_control, ux_state, ux_state.speed_index + 1)
return True
elif pressed(imgui.Key.Equal):
ux.set_speed_index(step_control, ux_state, ux_state.speed_index - 1)
return True
return False
def handle_camera_select_keyboard_events(
model: mujoco.MjModel,
camera: mujoco.MjvCamera,
ux_state: ux.UxState,
) -> bool:
"""Handles keyboard shortcuts for camera selection.
Args:
model: The MuJoCo model.
camera: The MuJoCo camera object.
ux_state: The UX state object.
Returns:
True if a key was handled, False otherwise.
"""
if imgui.GetIO().WantCaptureKeyboard:
return False
pressed = imgui.IsKeyChordPressed
if pressed(imgui.Key.Escape):
ux_state.camera_index = ux.set_camera(model, camera, ux.TUMBLE_CAMERA_IDX)
return True
elif pressed(imgui.Key.LeftBracket):
ux_state.camera_index = ux.set_camera(
model, camera, ux_state.camera_index - 1
)
return True
elif pressed(imgui.Key.RightBracket):
ux_state.camera_index = ux.set_camera(
model, camera, ux_state.camera_index + 1
)
return True
return False
def handle_freecam_wasd_keyboard_events(
model: mujoco.MjModel,
data: mujoco.MjData,
camera: mujoco.MjvCamera,
cam_speed: float,
) -> tuple[bool, float]:
"""Handles keyboard shortcuts for free camera movement.
Args:
model: The MuJoCo model.
data: The MuJoCo data.
camera: The MuJoCo camera object.
cam_speed: The current camera speed.
Returns:
A tuple of (handled, updated_cam_speed).
"""
if imgui.GetIO().WantCaptureKeyboard:
return False, cam_speed
if (
imgui.IsKeyDown(imgui.Key.W)
or imgui.IsKeyDown(imgui.Key.S)
or imgui.IsKeyDown(imgui.Key.A)
or imgui.IsKeyDown(imgui.Key.D)
or imgui.IsKeyDown(imgui.Key.Q)
or imgui.IsKeyDown(imgui.Key.E)
):
moved = False
# Move (dolly) forward/backward using W and S keys.
if imgui.IsKeyDown(imgui.Key.W):
ux.MoveCamera(
model,
data,
camera,
ux.CameraMotion.TRUCK_DOLLY,
0,
cam_speed,
)
moved = True
elif imgui.IsKeyDown(imgui.Key.S):
ux.MoveCamera(
model,
data,
camera,
ux.CameraMotion.TRUCK_DOLLY,
0,
-cam_speed,
)
moved = True
# Strafe (truck) left/right using A and D keys.
if imgui.IsKeyDown(imgui.Key.A):
ux.MoveCamera(
model,
data,
camera,
ux.CameraMotion.TRUCK_DOLLY,
-cam_speed,
0,
)
moved = True
elif imgui.IsKeyDown(imgui.Key.D):
ux.MoveCamera(
model,
data,
camera,
ux.CameraMotion.TRUCK_DOLLY,
cam_speed,
0,
)
moved = True
# Move (pedestal) up/down using Q and E keys.
if imgui.IsKeyDown(imgui.Key.Q):
ux.MoveCamera(
model,
data,
camera,
ux.CameraMotion.TRUCK_PEDESTAL,
0,
cam_speed,
)
moved = True
elif imgui.IsKeyDown(imgui.Key.E):
ux.MoveCamera(
model,
data,
camera,
ux.CameraMotion.TRUCK_PEDESTAL,
0,
-cam_speed,
)
moved = True
if moved:
cam_speed += 0.001
max_speed = 0.1 if imgui.GetIO().KeyShift else 0.01
if cam_speed > max_speed:
cam_speed = max_speed
else:
cam_speed = 0.001
return True, cam_speed
return False, cam_speed
def handle_keyboard_events(
model: mujoco.MjModel,
data: mujoco.MjData,
camera: mujoco.MjvCamera,
vis_options: mujoco.MjvOption,
step_control: sim.StepControl,
ux_state: ux.UxState,
cam_speed: float,
) -> tuple[bool, float]:
"""Handle keyboard events according to Studio's bindings.
Args:
model: The MuJoCo model.
data: The MuJoCo data.
camera: The MuJoCo camera object.
vis_options: The MuJoCo visualization options.
step_control: The simulation step control object.
ux_state: The UX state object.
cam_speed: The current camera speed.
Returns:
A tuple of (handled, updated_cam_speed).
"""
if imgui.GetIO().WantCaptureKeyboard:
return False, cam_speed
is_freecam_wasd = ux_state.camera_index == ux.FREE_CAMERA_IDX
if handle_step_control_keyboard_events(model, data, step_control, ux_state):
return True, cam_speed
if handle_camera_select_keyboard_events(model, camera, ux_state):
return True, cam_speed
if handle_vis_options_keyboard_events(vis_options, is_freecam_wasd):
return True, cam_speed
if is_freecam_wasd:
handled, cam_speed = handle_freecam_wasd_keyboard_events(
model, data, camera, cam_speed
)
if handled:
return True, cam_speed
return False, cam_speed
def handle_camera_tracking_mouse_events(
model: mujoco.MjModel,
data: mujoco.MjData,
camera: mujoco.MjvCamera,
vis_options: mujoco.MjvOption,
ux_state: ux.UxState,
) -> None:
"""Handles mouse events for camera tracking."""
io = imgui.GetIO()
if imgui.GetIO().WantCaptureMouse:
return
if io.DisplaySize.x <= 0 or io.DisplaySize.y <= 0:
return
mouse_x = io.MousePos.x / io.DisplaySize.x
mouse_y = io.MousePos.y / io.DisplaySize.y
aspect_ratio = io.DisplaySize.x / io.DisplaySize.y
# Right double click.
if imgui.IsMouseDoubleClicked(imgui.MouseButton.Right):
picked = ux.Pick(
model,
data,
camera,
mouse_x,
mouse_y,
aspect_ratio,
vis_options,
)
if picked.body > 0 and io.KeyCtrl:
# Switch camera to tracking mode and track the selected body.
camera.type = int(mujoco.mjtCamera.mjCAMERA_TRACKING)
camera.trackbodyid = picked.body
camera.fixedcamid = -1
ux_state.camera_index = ux.TRACKING_CAMERA_IDX
def handle_mouse_events(
model: mujoco.MjModel,
data: mujoco.MjData,
camera: mujoco.MjvCamera,
vis_options: mujoco.MjvOption,
perturb: mujoco.MjvPerturb,
ux_state: ux.UxState,
) -> None:
"""Handles mouse events."""
io = imgui.GetIO()
if io.WantCaptureMouse:
return
if io.DisplaySize.x <= 0 or io.DisplaySize.y <= 0:
return
mouse_x = io.MousePos.x / io.DisplaySize.x
mouse_y = io.MousePos.y / io.DisplaySize.y
mouse_dx = io.MouseDelta.x / io.DisplaySize.x
mouse_dy = io.MouseDelta.y / io.DisplaySize.y
mouse_scroll = io.MouseWheel / 50.0
is_mouse_moving = mouse_dx != 0.0 or mouse_dy != 0.0
is_any_mouse_down = (
imgui.IsMouseDown(imgui.MouseButton.Left)
or imgui.IsMouseDown(imgui.MouseButton.Right)
or imgui.IsMouseDown(imgui.MouseButton.Middle)
)
is_mouse_dragging = is_mouse_moving and is_any_mouse_down
# If no mouse buttons are down, end any active perturbations.
if not is_any_mouse_down:
perturb.active = 0
# Handle perturbation mouse actions.
if is_mouse_dragging and io.KeyCtrl:
if perturb.select > 0:
action = int(mujoco.mjtMouse.mjMOUSE_NONE)
if imgui.IsMouseDown(imgui.MouseButton.Left):
if io.KeyAlt:
action = int(
mujoco.mjtMouse.mjMOUSE_MOVE_H
if io.KeyShift
else mujoco.mjtMouse.mjMOUSE_MOVE_V
)
else:
action = int(
mujoco.mjtMouse.mjMOUSE_ROTATE_H
if io.KeyShift
else mujoco.mjtMouse.mjMOUSE_ROTATE_V
)
elif imgui.IsMouseDown(imgui.MouseButton.Right):
action = int(
mujoco.mjtMouse.mjMOUSE_MOVE_H
if io.KeyShift
else mujoco.mjtMouse.mjMOUSE_MOVE_V
)
elif imgui.IsMouseDown(imgui.MouseButton.Middle):
action = int(mujoco.mjtMouse.mjMOUSE_ZOOM)
active = int(
mujoco.mjtPertBit.mjPERT_TRANSLATE
if action
in (
int(mujoco.mjtMouse.mjMOUSE_MOVE_V),
int(mujoco.mjtMouse.mjMOUSE_MOVE_H),
)
else mujoco.mjtPertBit.mjPERT_ROTATE
)
if active != perturb.active:
ux.InitPerturb(model, data, camera, perturb, active)
ux.MovePerturb(
model,
data,
camera,
perturb,
action,
mouse_dx,
mouse_dy,
)
elif is_mouse_dragging:
if ux_state.camera_index == ux.FREE_CAMERA_IDX:
if imgui.IsMouseDown(imgui.MouseButton.Left):
ux.MoveCamera(
model,
data,
camera,
ux.CameraMotion.PAN_TILT,
mouse_dx,
mouse_dy,
)
else:
if imgui.IsMouseDown(imgui.MouseButton.Left):
ux.MoveCamera(
model,
data,
camera,
ux.CameraMotion.ORBIT,
mouse_dx,
mouse_dy,
)
elif imgui.IsMouseDown(imgui.MouseButton.Middle):
ux.MoveCamera(
model,
data,
camera,
ux.CameraMotion.ZOOM,
mouse_dx,
mouse_dy,
)
# Right mouse movement is relative to the horizontal and vertical planes.
if imgui.IsMouseDown(imgui.MouseButton.Right) and io.KeyShift:
ux.MoveCamera(
model,
data,
camera,
ux.CameraMotion.PLANAR_MOVE_H,
mouse_dx,
mouse_dy,
)
elif imgui.IsMouseDown(imgui.MouseButton.Right):
ux.MoveCamera(
model,
data,
camera,
ux.CameraMotion.PLANAR_MOVE_V,
mouse_dx,
mouse_dy,
)
# Mouse scroll.
if mouse_scroll != 0.0 and ux_state.camera_index != ux.FREE_CAMERA_IDX:
ux.MoveCamera(
model,
data,
camera,
ux.CameraMotion.ZOOM,
0,
-mouse_scroll,
)
aspect_ratio = io.DisplaySize.x / io.DisplaySize.y
# Left double click.
if imgui.IsMouseDoubleClicked(imgui.MouseButton.Left):
picked = ux.Pick(
model,
data,
camera,
mouse_x,
mouse_y,
aspect_ratio,
vis_options,
)
if picked.body >= 0:
perturb.select = picked.body
perturb.flexselect = picked.flex
perturb.skinselect = picked.skin
# Compute the local position of the selected object in the world.
tmp = np.array(picked.point, dtype=np.float64) - data.xpos[picked.body]
xmat = np.array(data.xmat[picked.body], dtype=np.float64).reshape(3, 3)
perturb.localpos = xmat.T @ tmp
else:
perturb.select = 0
perturb.flexselect = -1
perturb.skinselect = -1
handle_camera_tracking_mouse_events(
model, data, camera, vis_options, ux_state
)
python/mujoco/experimental/studio/ux.cc
+398
View File
@@ -0,0 +1,398 @@
// Copyright 2026 DeepMind Technologies Limited
//
// 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
//
// https://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.
// Python bindings for MuJoCo platform UX components.
#include <algorithm>
#include <array>
#include <string>
#include <tuple>
#include <vector>
#include <imgui.h>
#include <mujoco/mujoco.h>
#include "third_party/mujoco/src/experimental/platform/helpers.h"
#include "third_party/mujoco/src/experimental/platform/sim/step_control.h"
#include "third_party/mujoco/src/experimental/platform/ux/gui.h"
#include "third_party/mujoco/src/experimental/platform/ux/interaction.h"
#include "structs.h"
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
namespace py = pybind11;
struct UxState {
// Read/edited by step_control_gui
int speed_index = 0;
// Read/edited by state_gui
std::vector<mjtNum> state;
int state_sig = 0;
// Read/edited by watch_gui
char watch_field_name[256] = {0};
int watch_field_index = 0;
// Read/edited by noise_gui
float noise_scale = 0.0f;
float noise_rate = 0.0f;
// Read/edited by camera_selection_gui
int camera_index = mujoco::platform::kTumbleCameraIdx;
};
struct RenderFlags {
std::array<uint8_t, mjNRNDFLAG> flags = {0};
};
PYBIND11_MODULE(ux, m) {
py::class_<RenderFlags>(m, "RenderFlags")
.def(py::init<>())
.def_readwrite("flags", &RenderFlags::flags);
m.doc() = "MuJoCo platform UX components.";
py::enum_<mujoco::platform::GuiTheme>(m, "GuiTheme")
.value("LIGHT", mujoco::platform::GuiTheme::kLight)
.value("DARK", mujoco::platform::GuiTheme::kDark)
.value("CLASSIC", mujoco::platform::GuiTheme::kClassic);
py::class_<UxState>(m, "UxState")
.def(py::init<>())
.def_readwrite("speed_index", &UxState::speed_index)
.def_readwrite("state", &UxState::state)
.def_readwrite("state_sig", &UxState::state_sig)
.def_readwrite("watch_field_index", &UxState::watch_field_index)
.def_readwrite("noise_scale", &UxState::noise_scale)
.def_readwrite("noise_rate", &UxState::noise_rate)
.def_readwrite("camera_index", &UxState::camera_index)
.def_property(
"watch_field_name",
[](const UxState& self) {
return std::string(self.watch_field_name);
},
[](UxState& self, const std::string& val) {
std::snprintf(self.watch_field_name, sizeof(self.watch_field_name),
"%s", val.c_str());
});
m.def(
"setup_theme",
[](mujoco::platform::GuiTheme theme) {
mujoco::platform::SetupTheme(theme);
},
py::arg("theme"), "Set up Dear ImGui visual theme.");
m.def(
"configure_docking_layout",
[]() {
ImVec4 r = mujoco::platform::ConfigureDockingLayout();
return std::make_tuple(r.x, r.y, r.z, r.w);
},
"Configure the docking layout with Options (left) and Inspector (right) "
"panes. Returns (x, y, w, h) of the central workspace area.");
m.def(
"step_control_gui",
[](const mujoco::python::MjModelWrapper& model,
mujoco::platform::StepControl* step_control, UxState& ux_state) {
mujoco::platform::StepControlGui(model.get(), step_control,
ux_state.speed_index);
},
py::arg("model"), py::arg("step_control"), py::arg("ux_state"),
"Render the simulation stepping control GUI. Modifies "
"ux_state.speed_index.");
m.def(
"theme_select_gui",
[](mujoco::platform::GuiTheme theme) {
bool changed = mujoco::platform::ThemeSelectGui(&theme);
return std::make_tuple(changed, theme);
},
py::arg("theme"),
"Render the GUI theme selector. Returns (changed, theme).");
m.def(
"label_selection_gui",
[](mujoco::python::MjvOptionWrapper& vis_options) {
return mujoco::platform::LabelSelectionGui(vis_options.get());
},
py::arg("vis_options"), "Render the visualization label selection GUI.");
m.def(
"frame_selection_gui",
[](mujoco::python::MjvOptionWrapper& vis_options) {
return mujoco::platform::FrameSelectionGui(vis_options.get());
},
py::arg("vis_options"), "Render the visualization frame selection GUI.");
m.def(
"camera_selection_gui",
[](const mujoco::python::MjModelWrapper& model,
mujoco::python::MjDataWrapper& data,
mujoco::python::MjvCameraWrapper& camera, UxState& ux_state) {
bool changed = mujoco::platform::CameraSelectionGui(
model.get(), data.get(), *camera.get(), ux_state.camera_index);
return changed;
},
py::arg("model"), py::arg("data"), py::arg("camera"), py::arg("ux_state"),
"Render the camera selection GUI. Modifies ux_state.camera_index. "
"Returns true if camera changed.");
m.def(
"set_camera",
[](const mujoco::python::MjModelWrapper& model,
mujoco::python::MjvCameraWrapper& camera, int request_idx) {
return mujoco::platform::SetCamera(model.get(), camera.get(), request_idx);
},
py::arg("model"), py::arg("camera"), py::arg("request_idx"),
"Set the camera index and update the camera object.");
m.def(
"set_speed_index",
[](mujoco::platform::StepControl* step_control, UxState& ux_state, int idx) {
mujoco::platform::SetSpeedIndex(step_control, ux_state.speed_index, idx);
},
py::arg("step_control"), py::arg("ux_state"), py::arg("idx"),
"Set the simulation speed index.");
m.def(
"physics_gui",
[](mujoco::python::MjModelWrapper& model, float min_width) {
mujoco::platform::PhysicsGui(model.get(), min_width);
},
py::arg("model"), py::arg("min_width") = 150.0f,
"Render the physics settings UI.");
m.def(
"rendering_gui",
[](const mujoco::python::MjModelWrapper& model,
mujoco::python::MjvOptionWrapper& vis_options,
RenderFlags& render_flags) {
mjtByte flags[mjNRNDFLAG] = {0};
for (int i = 0; i < mjNRNDFLAG; ++i) {
flags[i] = render_flags.flags[i];
}
mujoco::platform::RenderingGui(model.get(), vis_options.get(), flags,
150.0f);
for (int i = 0; i < mjNRNDFLAG; ++i) {
render_flags.flags[i] = flags[i];
}
},
py::arg("model"), py::arg("vis_options"), py::arg("render_flags"),
"Render the rendering settings UI. Modifies render_flags in place.");
m.def(
"groups_gui",
[](const mujoco::python::MjModelWrapper& model,
mujoco::python::MjvOptionWrapper& vis_options, float min_width) {
mujoco::platform::GroupsGui(model.get(), vis_options.get(), min_width);
},
py::arg("model"), py::arg("vis_options"), py::arg("min_width") = 150.0f,
"Render the visibility groups UI.");
m.def(
"visualization_gui",
[](mujoco::python::MjModelWrapper& model,
mujoco::python::MjvOptionWrapper& vis_options,
mujoco::python::MjvCameraWrapper& camera, float min_width) {
mujoco::platform::VisualizationGui(model.get(), vis_options.get(),
camera.get(), min_width);
},
py::arg("model"), py::arg("vis_options"), py::arg("camera"),
py::arg("min_width") = 150.0f, "Render the visualization settings UI.");
m.def(
"controls_gui",
[](const mujoco::python::MjModelWrapper& model,
mujoco::python::MjDataWrapper& data,
mujoco::python::MjvOptionWrapper& vis_options) {
mujoco::platform::ControlsGui(model.get(), data.get(),
vis_options.get());
},
py::arg("model"), py::arg("data"), py::arg("vis_options"),
"Render the actuator controls UI.");
m.def(
"joints_gui",
[](const mujoco::python::MjModelWrapper& model,
mujoco::python::MjDataWrapper& data,
mujoco::python::MjvOptionWrapper& vis_options) {
mujoco::platform::JointsGui(model.get(), data.get(), vis_options.get());
},
py::arg("model"), py::arg("data"), py::arg("vis_options"),
"Render the joints UI.");
m.def(
"sensor_gui",
[](const mujoco::python::MjModelWrapper& model,
mujoco::python::MjDataWrapper& data) {
mujoco::platform::SensorGui(model.get(), data.get());
},
py::arg("model"), py::arg("data"), "Render the sensor data plot.");
m.def(
"state_gui",
[](const mujoco::python::MjModelWrapper& model,
mujoco::python::MjDataWrapper& data, UxState& ux_state,
float min_width) {
mujoco::platform::StateGui(model.get(), data.get(), ux_state.state,
ux_state.state_sig, min_width);
},
py::arg("model"), py::arg("data"), py::arg("ux_state"),
py::arg("min_width") = 150.0f,
"Render the state UI. Modifies ux_state.state and ux_state.state_sig.");
m.def(
"watch_gui",
[](const mujoco::python::MjModelWrapper& model,
mujoco::python::MjDataWrapper& data, UxState& ux_state) {
mujoco::platform::WatchGui(
model.get(), data.get(), ux_state.watch_field_name,
sizeof(ux_state.watch_field_name), ux_state.watch_field_index);
},
py::arg("model"), py::arg("data"), py::arg("ux_state"),
"Render the watch UI. Modifies ux_state.watch_field_name and "
"ux_state.watch_field_index.");
m.def(
"noise_gui",
[](const mujoco::python::MjModelWrapper& model,
mujoco::python::MjDataWrapper& data, UxState& ux_state) {
mujoco::platform::NoiseGui(model.get(), data.get(),
ux_state.noise_scale, ux_state.noise_rate);
},
py::arg("model"), py::arg("data"), py::arg("ux_state"),
"Render the noise UI. Modifies ux_state.noise_scale and "
"ux_state.noise_rate.");
m.def(
"convergence_gui",
[](const mujoco::python::MjModelWrapper& model,
mujoco::python::MjDataWrapper& data) {
mujoco::platform::ConvergenceGui(model.get(), data.get());
},
py::arg("model"), py::arg("data"),
"Render the solver convergence chart.");
m.def(
"counts_gui",
[](const mujoco::python::MjModelWrapper& model,
mujoco::python::MjDataWrapper& data) {
mujoco::platform::CountsGui(model.get(), data.get());
},
py::arg("model"), py::arg("data"), "Render the solver counts chart.");
m.def(
"stats_gui",
[](const mujoco::python::MjModelWrapper& model,
mujoco::python::MjDataWrapper& data, bool paused, float fps) {
mujoco::platform::StatsGui(model.get(), data.get(), paused, fps);
},
py::arg("model"), py::arg("data"), py::arg("paused"), py::arg("fps"),
"Render the simulation statistics UI.");
m.attr("FREE_CAMERA_IDX") = mujoco::platform::kFreeCameraIdx;
m.attr("TUMBLE_CAMERA_IDX") = mujoco::platform::kTumbleCameraIdx;
m.attr("TRACKING_CAMERA_IDX") = mujoco::platform::kTrackingCameraIdx;
py::enum_<mujoco::platform::CameraMotion>(m, "CameraMotion")
.value("ZOOM", mujoco::platform::CameraMotion::ZOOM)
.value("ORBIT", mujoco::platform::CameraMotion::ORBIT)
.value("TRUCK_PEDESTAL", mujoco::platform::CameraMotion::TRUCK_PEDESTAL)
.value("TRUCK_DOLLY", mujoco::platform::CameraMotion::TRUCK_DOLLY)
.value("PAN_TILT", mujoco::platform::CameraMotion::PAN_TILT)
.value("PLANAR_MOVE_H", mujoco::platform::CameraMotion::PLANAR_MOVE_H)
.value("PLANAR_MOVE_V", mujoco::platform::CameraMotion::PLANAR_MOVE_V)
.export_values();
m.def(
"MoveCamera",
[](const mujoco::python::MjModelWrapper& model,
const mujoco::python::MjDataWrapper& data,
mujoco::python::MjvCameraWrapper& cam,
mujoco::platform::CameraMotion motion, mjtNum dx, mjtNum dy) {
mujoco::platform::MoveCamera(model.get(), data.get(), cam.get(), motion,
dx, dy);
},
py::arg("model"), py::arg("data"), py::arg("cam"), py::arg("motion"),
py::arg("dx"), py::arg("dy"), "Moves the given camera.");
m.def(
"InitPerturb",
[](const mujoco::python::MjModelWrapper& model,
const mujoco::python::MjDataWrapper& data,
const mujoco::python::MjvCameraWrapper& cam,
mujoco::python::MjvPerturbWrapper& pert, int active) {
mujoco::platform::InitPerturb(model.get(), data.get(), cam.get(),
pert.get(),
static_cast<mjtPertBit>(active));
},
py::arg("model"), py::arg("data"), py::arg("cam"), py::arg("pert"),
py::arg("active"), "Initializes mouse perturbation.");
m.def(
"MovePerturb",
[](const mujoco::python::MjModelWrapper& model,
const mujoco::python::MjDataWrapper& data,
const mujoco::python::MjvCameraWrapper& cam,
mujoco::python::MjvPerturbWrapper& pert, int action, mjtNum reldx,
mjtNum reldy) {
mujoco::platform::MovePerturb(model.get(), data.get(), cam.get(),
pert.get(), static_cast<mjtMouse>(action),
reldx, reldy);
},
py::arg("model"), py::arg("data"), py::arg("cam"), py::arg("pert"),
py::arg("action"), py::arg("reldx"), py::arg("reldy"),
"Moves mouse perturbation.");
py::class_<mujoco::platform::PickResult>(m, "PickResult")
.def_readwrite("dist", &mujoco::platform::PickResult::dist)
.def_readwrite("body", &mujoco::platform::PickResult::body)
.def_readwrite("geom", &mujoco::platform::PickResult::geom)
.def_readwrite("flex", &mujoco::platform::PickResult::flex)
.def_readwrite("skin", &mujoco::platform::PickResult::skin)
.def_property(
"point",
[](const mujoco::platform::PickResult& res) {
return py::make_tuple(res.point[0], res.point[1], res.point[2]);
},
[](mujoco::platform::PickResult& res, const py::tuple& t) {
res.point[0] = t[0].cast<mjtNum>();
res.point[1] = t[1].cast<mjtNum>();
res.point[2] = t[2].cast<mjtNum>();
});
m.def(
"Pick",
[](const mujoco::python::MjModelWrapper& model,
const mujoco::python::MjDataWrapper& data,
const mujoco::python::MjvCameraWrapper& cam, float x, float y,
float aspect_ratio, const mujoco::python::MjvOptionWrapper& opt) {
return mujoco::platform::Pick(model.get(), data.get(), cam.get(), x, y,
aspect_ratio, opt.get());
},
py::arg("model"), py::arg("data"), py::arg("cam"), py::arg("x"),
py::arg("y"), py::arg("aspect_ratio"), py::arg("opt"),
"Picks object under cursor.");
m.def(
"camera_to_string",
[](const mujoco::python::MjDataWrapper& data,
const mujoco::python::MjvCameraWrapper& camera) {
return mujoco::platform::CameraToString(data.get(), camera.get());
},
py::arg("data"), py::arg("camera"),
"Returns an XML string representation of the camera.");
}
python/mujoco/experimental/studio/viewer_protocol.py
+43
View File
@@ -0,0 +1,43 @@
# Copyright 2026 DeepMind Technologies Limited
#
# 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
#
# https://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.
"""Structural protocol defining the common viewer interface.
StudioApp uses the protocol for convenience methods that accept any viewer.
"""
from typing import Protocol
import mujoco
from mujoco.experimental.studio import ux
class Viewer(Protocol):
"""Structural interface for any viewer."""
camera: mujoco.MjvCamera
perturb: mujoco.MjvPerturb
vis_options: mujoco.MjvOption
render_flags: ux.RenderFlags
def is_running(self) -> bool:
...
def sync(self, model: mujoco.MjModel, data: mujoco.MjData) -> None:
...
def stop(self) -> None:
...
def get_drop_file(self) -> str:
...
python/mujoco/functions.cc
+3 -6
View File
@@ -591,15 +591,12 @@ PYBIND11_MODULE(_functions, pymodule) {
Def<traits::mj_id2name>(pymodule);
Def<traits::mj_fullM>(
pymodule,
[](const raw::MjModel* m, Eigen::Ref<EigenArrayXX> dst,
Eigen::Ref<const EigenVectorX> M) {
if (M.size() != m->nM) {
throw py::type_error("M should be of size nM");
}
[](const raw::MjModel* m, const raw::MjData* d,
Eigen::Ref<EigenArrayXX> dst) {
if (dst.cols() != m->nv || dst.rows() != m->nv) {
throw py::type_error("dst should be of shape (nv, nv)");
}
return ::mj_fullM(m, dst.data(), M.data());
return ::mj_fullM(m, d, dst.data());
});
Def<traits::mj_mulM>(
pymodule,
python/mujoco/introspect/functions.py
+7 -83
View File
@@ -3478,20 +3478,20 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
inner_type=ValueType(name='mjModel', is_const=True),
),
),
FunctionParameterDecl(
name='d',
type=PointerType(
inner_type=ValueType(name='mjData', is_const=True),
),
),
FunctionParameterDecl(
name='dst',
type=PointerType(
inner_type=ValueType(name='mjtNum'),
),
),
FunctionParameterDecl(
name='M',
type=PointerType(
inner_type=ValueType(name='mjtNum', is_const=True),
),
),
),
doc='Convert sparse inertia matrix M into full (i.e. dense) matrix.',
doc='Convert sparse inertia matrix into full (i.e. dense) matrix.',
)),
('mj_mulM',
FunctionDecl(
@@ -6288,44 +6288,6 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
),
doc='Main error function; does not return to caller.',
)),
('mju_error_i',
FunctionDecl(
name='mju_error_i',
return_type=ValueType(name='void'),
parameters=(
FunctionParameterDecl(
name='msg',
type=PointerType(
inner_type=ValueType(name='char', is_const=True),
),
),
FunctionParameterDecl(
name='i',
type=ValueType(name='int'),
),
),
doc='Deprecated: use mju_error.',
)),
('mju_error_s',
FunctionDecl(
name='mju_error_s',
return_type=ValueType(name='void'),
parameters=(
FunctionParameterDecl(
name='msg',
type=PointerType(
inner_type=ValueType(name='char', is_const=True),
),
),
FunctionParameterDecl(
name='text',
type=PointerType(
inner_type=ValueType(name='char', is_const=True),
),
),
),
doc='Deprecated: use mju_error.',
)),
('mju_warning',
FunctionDecl(
name='mju_warning',
@@ -6340,44 +6302,6 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
),
doc='Main warning function; returns to caller.',
)),
('mju_warning_i',
FunctionDecl(
name='mju_warning_i',
return_type=ValueType(name='void'),
parameters=(
FunctionParameterDecl(
name='msg',
type=PointerType(
inner_type=ValueType(name='char', is_const=True),
),
),
FunctionParameterDecl(
name='i',
type=ValueType(name='int'),
),
),
doc='Deprecated: use mju_warning.',
)),
('mju_warning_s',
FunctionDecl(
name='mju_warning_s',
return_type=ValueType(name='void'),
parameters=(
FunctionParameterDecl(
name='msg',
type=PointerType(
inner_type=ValueType(name='char', is_const=True),
),
),
FunctionParameterDecl(
name='text',
type=PointerType(
inner_type=ValueType(name='char', is_const=True),
),
),
),
doc='Deprecated: use mju_warning.',
)),
('mju_clearHandlers',
FunctionDecl(
name='mju_clearHandlers',
src/engine/engine_collision_gjk.c
+17 -8
View File
@@ -28,6 +28,11 @@
#define mjMINVAL2 (mjMINVAL * mjMINVAL)
#define mjMAXVAL2 (mjMAXVAL * mjMAXVAL)
// align memory size on 8-byte boundary; needed for single precision
static inline size_t align8(size_t size) {
return ((size + 7) / 8) * 8;
}
// subdistance algorithm for GJK that computes the barycentric coordinates of the point in a
// simplex closest to the origin
// implementation adapted from Montanari et al, ToG 2017
@@ -1533,6 +1538,9 @@ static mjtNum planeNormal(mjtNum res[3], const mjtNum v1[3], const mjtNum v2[3],
sub3(diff1, v2, v1);
sub3(diff2, v3, v1);
cross3(res, diff1, diff2);
// normalize isn't needed (cancelled out), but done to avoid asymmetric rounding later on
mju_normalize3(res);
return dot3(res, v1);
}
@@ -2217,10 +2225,11 @@ static inline void inflate(mjCCDStatus* status, mjtNum margin1, mjtNum margin2)
// return size in bytes of the buffer needed for mjc_ccd for a given number of iterations
size_t mjc_ccdSize(int iterations) {
return (sizeof(Face) * 6 * iterations) // faces in polytope
+ (sizeof(Face*) * 6 * iterations) // map in polytope
+ (sizeof(Vertex) * (5 + iterations)) // vertices in polytope
+ 2 * (24 * sizeof(int)); // horizon data
return align8(sizeof(Vertex) * (5 + iterations)) // vertices in polytope
+ align8(sizeof(Face) * 6 * iterations) // faces in polytope
+ align8(sizeof(Face*) * 6 * iterations) // map in polytope
+ align8(sizeof(int) * 24) // horizon indices
+ align8(sizeof(int) * 24); // horizon edges
}
@@ -2313,13 +2322,13 @@ mjtNum mjc_ccd(const mjCCDConfig* config, mjCCDStatus* status, mjCCDObj* obj1, m
pt.maxfaces = 6 * N;
uint8_t* buffer = config->buffer;
pt.verts = (Vertex*)buffer;
buffer += sizeof(Vertex) * (5 + N);
buffer += align8(sizeof(Vertex) * (5 + N));
pt.faces = (Face*)buffer;
buffer += sizeof(Face) * (6 * N);
buffer += align8(sizeof(Face) * (6 * N));
pt.map = (Face**)buffer;
buffer += sizeof(Face*) * (6 * N);
buffer += align8(sizeof(Face*) * (6 * N));
pt.horizon.indices = (int*)buffer;
buffer += sizeof(int) * 24;
buffer += align8(sizeof(int) * 24);
pt.horizon.edges = (int*)buffer;
int ret;
src/engine/engine_support.c
+2 -13
View File
@@ -367,19 +367,8 @@ void mj_setKeyframe(mjModel* m, const mjData* d, int k) {
//-------------------------- inertia functions -----------------------------------------------------
// convert sparse inertia matrix M into full matrix
void mj_fullM(const mjModel* m, mjtNum* dst, const mjtNum* M) {
int adr = 0, nv = m->nv;
mju_zero(dst, nv*nv);
for (int i=0; i < nv; i++) {
int j = i;
while (j >= 0) {
dst[i*nv+j] = M[adr];
dst[j*nv+i] = M[adr];
j = m->dof_parentid[j];
adr++;
}
}
void mj_fullM(const mjModel* m, const mjData* d, mjtNum* dst) {
mju_sym2dense(dst, d->M, m->nv, m->M_rownnz, m->M_rowadr, m->M_colind);
}
src/engine/engine_support.h
+1 -1
View File
@@ -58,7 +58,7 @@ MJAPI void mj_setKeyframe(mjModel* m, const mjData* d, int k);
//-------------------------- inertia functions -----------------------------------------------------
// convert sparse inertia matrix M into full matrix
MJAPI void mj_fullM(const mjModel* m, mjtNum* dst, const mjtNum* M);
MJAPI void mj_fullM(const mjModel* m, const mjData* d, mjtNum* dst);
// multiply vector by inertia matrix
MJAPI void mj_mulM(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec);
src/engine/engine_util_errmem.c
-24
View File
@@ -169,30 +169,6 @@ void mju_warning(const char* msg, ...) {
}
// error with int argument
void mju_error_i(const char* msg, int i) {
mju_error(msg, i);
}
// warning with int argument
void mju_warning_i(const char* msg, int i) {
mju_warning(msg, i);
}
// error string argument
void mju_error_s(const char* msg, const char* text) {
mju_error(msg, text);
}
// warning string argument
void mju_warning_s(const char* msg, const char* text) {
mju_warning(msg, text);
}
//------------------------------ malloc and free ---------------------------------------------------
// allocate memory; byte-align on 64; pad size to multiple of 64
src/engine/engine_util_errmem.h
-4
View File
@@ -58,13 +58,9 @@ MJAPI void _mjPRIVATE__set_tls_warning_fn(void (*h)(const char*));
MJAPI void mju_error_raw(const char* msg);
MJAPI void mju_error(const char* msg, ...) mjPRINTFLIKE(1, 2);
MJAPI void mju_error_v(const char* msg, va_list args);
MJAPI void mju_error_i(const char* msg, int i);
MJAPI void mju_error_s(const char* msg, const char* text);
// warnings
MJAPI void mju_warning(const char* msg, ...) mjPRINTFLIKE(1, 2);
MJAPI void mju_warning_i(const char* msg, int i);
MJAPI void mju_warning_s(const char* msg, const char* text);
// write [datetime, type: message] to MUJOCO_LOG.TXT
MJAPI void mju_writeLog(const char* type, const char* msg);
src/experimental/filament/compat/scene_geom_util.cc
+1 -2
View File
@@ -159,8 +159,6 @@ static void UpdateGeomMaterial(mjrRenderable* renderable, const mjvGeom& geom,
material.color[2] = geom.rgba[2];
material.color[3] = geom.rgba[3];
mjrf_setRenderableLayerMask(renderable, geom.category);
if (geom.matid >= 0 && geom.matid < model->nmat) {
auto get_texture = [&](int role) -> const mjrTexture* {
const int tex_id = model->mat_texid[geom.matid * mjNTEXROLE + role];
@@ -283,6 +281,7 @@ UniquePtr<mjrRenderable> CreateGeomRenderable(
const mjtByte render_flags[mjNRNDFLAG]) {
mjrRenderableParams params;
mjr_defaultRenderableParams(&params);
params.layer_mask = geom.category;
auto renderable = CreateRenderable(ctx, params);
PrepareGeomMeshes(renderable.get(), geom, model_objs);
UpdateGeomMaterial(renderable.get(), geom, model_objs, render_flags);
src/experimental/filament/filament/scene_view.cc
+2
View File
@@ -255,7 +255,9 @@ void SceneView::Render(filament::Renderer* renderer, const mjrRenderRequest& req
for (auto& iter : renderables_) {
iter->BindMaterialInstance(request);
}
for (auto& iter : renderables_) {
if (RenderTarget* target = iter->GetReflectionTarget()) {
viewport.left = 0;
viewport.bottom = 0;
src/experimental/filament/render_context_filament.cc
+5 -15
View File
@@ -245,6 +245,11 @@ void mjrf_setRenderableMaterial(mjrRenderable* renderable,
mujoco::Renderable::downcast(renderable)->UpdateMaterial(*material);
}
void mjrf_getRenderableMaterial(mjrRenderable* renderable,
mjrMaterial* material) {
*material = mujoco::Renderable::downcast(renderable)->GetMaterial();
}
void mjrf_setRenderableTransform(mjrRenderable* renderable,
const float position[3],
const float rotation[9]) {
@@ -261,21 +266,6 @@ void mjrf_setRenderableSize(mjrRenderable* renderable, const float size[3]) {
mujoco::Renderable::downcast(renderable)->SetSize(fsize);
}
void mjrf_setRenderableLayerMask(mjrRenderable* renderable,
uint8_t layer_mask) {
mujoco::Renderable::downcast(renderable)->SetLayerMask(layer_mask);
}
void mjrf_setRenderableCastShadows(mjrRenderable* renderable,
mjtByte cast_shadows) {
mujoco::Renderable::downcast(renderable)->SetCastShadows(cast_shadows);
}
void mjrf_setRenderableReceiveShadows(mjrRenderable* renderable,
mjtByte receive_shadows) {
mujoco::Renderable::downcast(renderable)->SetReceiveShadows(receive_shadows);
}
void mjrf_addLightToScene(mjrScene* scene, mjrLight* light) {
mujoco::SceneView::downcast(scene)->AddToScene(
mujoco::Light::downcast(light));
src/experimental/filament/render_context_filament.h
+20 -12
View File
@@ -67,7 +67,6 @@ struct mjrLight {};
struct mjrRenderable {};
struct mjrRenderTarget {};
// ## Rendering Context (mjrfContext)
//
// The Context is the main entry point for the library. It manages all the
@@ -520,22 +519,31 @@ typedef mjtLightType mjrLightType;
struct mjrLightParams {
// The type of light (e.g. spot, point, directional, etc.)
mjrLightType type;
// The texture to use for image lights.
const mjrTexture* texture;
// The color of the light.
float color[3];
// The intensity of the light, in candela.
float intensity;
// Whether or not the light casts shadows.
mjtByte cast_shadows;
// The range/distance in which the light is effective, in meters.
float range;
// The angle of the spot light cone, in degrees.
float spot_cone_angle;
// The radius of the bulb used for soft shadows.
float bulb_radius;
// The size of the shadow map.
int shadow_map_size;
// Blur width for EL VSM.
float vsm_blur_width;
};
@@ -656,15 +664,19 @@ void mjr_defaultMaterial(mjrMaterial* material);
struct mjrRenderableParams {
// Whether or not the Renderable casts shadows.
mjtByte cast_shadows;
// Whether or not the Renderable receives shadows.
mjtByte receive_shadows;
// The layers to which the Renderable belongs. This mask is used in
// conjunction with the layer mask in the Scene to determine which
// Renderables to render. Defaults to 0xff.
uint8_t layer_mask;
// Controls the order in which the Renderable is drawn relative to other
// Renderables; defaults to 4.
uint8_t priority;
// Similar to priority, but provides finer-grained control for Renderables
// with transparency; defaults to 0.
uint16_t blend_order;
@@ -694,6 +706,10 @@ void mjrf_setRenderableGeomMesh(mjrRenderable* renderable, mjtGeom type,
void mjrf_setRenderableMaterial(mjrRenderable* renderable,
const mjrMaterial* material);
// Copies the material properties of the renderable into the given mjrMaterial.
void mjrf_getRenderableMaterial(mjrRenderable* renderable,
mjrMaterial* material);
// Sets the transform position and rotation of the renderable.
void mjrf_setRenderableTransform(mjrRenderable* renderable,
const float position[3],
@@ -705,17 +721,6 @@ void mjrf_setRenderableTransform(mjrRenderable* renderable,
// capsule are scaled such that they always remain spherical).
void mjrf_setRenderableSize(mjrRenderable* renderable, const float size[3]);
// Sets whether the renderable casts shadows or not.
void mjrf_setRenderableCastShadows(mjrRenderable* renderable,
mjtByte cast_shadows);
// Sets whether the renderable receives shadows or not.
void mjrf_setRenderableReceiveShadows(mjrRenderable* renderable,
mjtByte receive_shadows);
// Sets the layer mask of the renderable. See mjrRenderableParams for details.
void mjrf_setRenderableLayerMask(mjrRenderable* renderable, uint8_t layer_mask);
// ## Render Targets (mjrRenderTarget)
//
// A render target is a memory buffer that holds the results of a rendering
@@ -726,10 +731,13 @@ void mjrf_setRenderableLayerMask(mjrRenderable* renderable, uint8_t layer_mask);
struct mjrRenderTargetConfig {
// The width of the render target.
int width;
// The height of the render target.
int height;
// The format of the color buffer in the render target.
mjrPixelFormat color_format;
// The format of the depth buffer in the render target.
mjrPixelFormat depth_format;
};
src/experimental/platform/sim/step_control.cc
+8 -2
View File
@@ -17,6 +17,7 @@
#include <algorithm>
#include <chrono>
#include <cstdlib>
#include <functional>
#include <optional>
#include <ratio>
@@ -91,7 +92,8 @@ StepControl::PauseState StepControl::GetPauseState() const {
return pause_state_;
}
StepControl::Status StepControl::Advance(mjModel* m, mjData* d) {
StepControl::Status StepControl::Advance(mjModel* m, mjData* d,
StepFn step_fn) {
if (!m) {
return Status::kOk;
}
@@ -182,7 +184,11 @@ StepControl::Status StepControl::Advance(mjModel* m, mjData* d) {
mjtNum prev_time = d->time;
InjectNoise(m, d);
mj_step(m, d);
if (step_fn) {
step_fn(m, d);
} else {
mj_step(m, d);
}
if (mjDISABLED(mjDSBL_AUTORESET)) {
for (mjtWarning w : kDivergedWarnings) {
src/experimental/platform/sim/step_control.h
+3 -1
View File
@@ -16,6 +16,7 @@
#define MUJOCO_SRC_EXPERIMENTAL_PLATFORM_SIM_STEP_CONTROL_H_
#include <chrono>
#include <functional>
#include <string>
#include <mujoco/mujoco.h>
@@ -24,6 +25,7 @@ namespace mujoco::platform {
using Seconds = std::chrono::duration<double>;
using Clock = std::chrono::steady_clock;
using StepFn = std::function<void(mjModel*, mjData*)>;
// State and logic for physics synchronization and stepping.
class StepControl {
@@ -53,7 +55,7 @@ class StepControl {
mjWARN_BADQACC, mjWARN_BADQVEL, mjWARN_BADQPOS};
// Steps physics forward, respecting speed settings and refresh budget.
Status Advance(mjModel* m, mjData* d);
Status Advance(mjModel* m, mjData* d, StepFn step_fn = nullptr);
// Ensures the next call to Advance() will synchronize time and step once.
void ForceSync();
src/experimental/studio/live.js
+8
View File
@@ -50,6 +50,14 @@ function resolveScheme(url) {
}
async function prefetchModelAssets(rootUrl, onProgress) {
// Only prefetch when the root URL is one we know how to fetch from
// Javascript. Other resource-provider-backed schemes (e.g., uploaded files
// via the drag-and-drop path) skip the prefetcher entirely and let
// Module.loadUrl handle take the existing slow path.
if (!/^(https?:|github:)/.test(rootUrl)) {
return { files: 0, bytes: 0, errors: 0 };
}
const primed = new Set(); // URLs we've already pushed into FetchCache
const stats = { files: 0, bytes: 0, errors: 0 };
src/experimental/usd/CMakeLists.txt
+40
View File
@@ -199,6 +199,42 @@ target_link_libraries(mujoco PUBLIC
${MJC_PHYSICS_PLUGIN_TARGET_NAME}
)
## ----- Newton USD Schemas (codeless plugin) -----
function(install_newton_usd_plugin install_base_dir)
include(FetchContent)
FetchContent_Declare(
newton-usd-schemas
GIT_REPOSITORY https://github.com/newton-physics/newton-usd-schemas.git
GIT_TAG v0.1.0rc3
GIT_SHALLOW TRUE
UPDATE_DISCONNECTED TRUE
)
FetchContent_GetProperties(newton-usd-schemas)
if(NOT newton-usd-schemas_POPULATED)
FetchContent_Populate(newton-usd-schemas)
endif()
set(NEWTON_USD_DIR "${newton-usd-schemas_SOURCE_DIR}/newton_usd_schemas")
if(NOT EXISTS "${NEWTON_USD_DIR}")
message(FATAL_ERROR "newton_usd_schemas directory not found in fetched repository")
endif()
set(NEWTON_BUILD_DIR "${CMAKE_BINARY_DIR}/${install_base_dir}/newton")
file(MAKE_DIRECTORY "${NEWTON_BUILD_DIR}")
configure_file("${NEWTON_USD_DIR}/plugInfo.json" "${NEWTON_BUILD_DIR}/plugInfo.json" COPYONLY)
configure_file("${NEWTON_USD_DIR}/generatedSchema.usda" "${NEWTON_BUILD_DIR}/generatedSchema.usda" COPYONLY)
install(FILES "${NEWTON_BUILD_DIR}/plugInfo.json"
DESTINATION "${install_base_dir}/newton"
)
install(FILES "${NEWTON_BUILD_DIR}/generatedSchema.usda"
DESTINATION "${install_base_dir}/newton"
)
endfunction()
## Installation
# Generate and install plugInfo.json for each plugin
@@ -219,6 +255,10 @@ install(FILES
DESTINATION ${MJ_USD_INSTALL_DIR_LIB}/mjcPhysics
)
install_newton_usd_plugin(
${MJ_USD_INSTALL_DIR_LIB}
)
# Install shared libraries
install(TARGETS
${MJCF_PLUGIN_TARGET_NAME}
src/experimental/usd/mjcPhysics/equalityJointAPI.cpp
+34 -2
View File
@@ -79,6 +79,28 @@ const TfType& MjcPhysicsEqualityJointAPI::_GetTfType() const {
return _GetStaticTfType();
}
UsdAttribute MjcPhysicsEqualityJointAPI::GetSolImpAttr() const {
return GetPrim().GetAttribute(MjcPhysicsTokens->mjcSolimp);
}
UsdAttribute MjcPhysicsEqualityJointAPI::CreateSolImpAttr(
VtValue const& defaultValue, bool writeSparsely) const {
return UsdSchemaBase::_CreateAttr(
MjcPhysicsTokens->mjcSolimp, SdfValueTypeNames->DoubleArray,
/* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely);
}
UsdAttribute MjcPhysicsEqualityJointAPI::GetSolRefAttr() const {
return GetPrim().GetAttribute(MjcPhysicsTokens->mjcSolref);
}
UsdAttribute MjcPhysicsEqualityJointAPI::CreateSolRefAttr(
VtValue const& defaultValue, bool writeSparsely) const {
return UsdSchemaBase::_CreateAttr(
MjcPhysicsTokens->mjcSolref, SdfValueTypeNames->DoubleArray,
/* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely);
}
UsdAttribute MjcPhysicsEqualityJointAPI::GetCoef0Attr() const {
return GetPrim().GetAttribute(MjcPhysicsTokens->mjcCoef0);
}
@@ -134,6 +156,15 @@ UsdAttribute MjcPhysicsEqualityJointAPI::CreateCoef4Attr(
/* custom = */ false, SdfVariabilityUniform, defaultValue, writeSparsely);
}
UsdRelationship MjcPhysicsEqualityJointAPI::GetMjcTargetRel() const {
return GetPrim().GetRelationship(MjcPhysicsTokens->mjcTarget);
}
UsdRelationship MjcPhysicsEqualityJointAPI::CreateMjcTargetRel() const {
return GetPrim().CreateRelationship(MjcPhysicsTokens->mjcTarget,
/* custom = */ false);
}
namespace {
static inline TfTokenVector _ConcatenateAttributeNames(
const TfTokenVector& left, const TfTokenVector& right) {
@@ -149,8 +180,9 @@ static inline TfTokenVector _ConcatenateAttributeNames(
const TfTokenVector& MjcPhysicsEqualityJointAPI::GetSchemaAttributeNames(
bool includeInherited) {
static TfTokenVector localNames = {
MjcPhysicsTokens->mjcCoef0, MjcPhysicsTokens->mjcCoef1,
MjcPhysicsTokens->mjcCoef2, MjcPhysicsTokens->mjcCoef3,
MjcPhysicsTokens->mjcSolimp, MjcPhysicsTokens->mjcSolref,
MjcPhysicsTokens->mjcCoef0, MjcPhysicsTokens->mjcCoef1,
MjcPhysicsTokens->mjcCoef2, MjcPhysicsTokens->mjcCoef3,
MjcPhysicsTokens->mjcCoef4,
};
static TfTokenVector allNames = _ConcatenateAttributeNames(
src/experimental/usd/mjcPhysics/generatedSchema.usda
+65 -17
View File
@@ -4,6 +4,7 @@
)
class "MjcSceneAPI" (
apiSchemas = ["NewtonSceneAPI"]
doc = "API providing global simulation options for MuJoCo."
)
{
@@ -120,7 +121,9 @@ class "MjcSceneAPI" (
)
uniform bool mjc:flag:gravity = 1 (
displayName = "Gravity Toggle"
doc = "Enables the application of gravitational acceleration as defined in mjOption."
doc = """DEPRECATED: Use newton:gravityEnabled instead.
Enables the application of gravitational acceleration as defined in mjOption."""
)
uniform bool mjc:flag:invdiscrete = 0 (
displayName = "Discrete-Time Inverse Dynamics Toggle"
@@ -200,7 +203,9 @@ class "MjcSceneAPI" (
)
uniform int mjc:option:iterations = 100 (
displayName = "Solver Iterations"
doc = "Maximum number of iterations of the constraint solver."
doc = """DEPRECATED: Use newton:maxSolverIterations instead.
Maximum number of iterations of the constraint solver."""
)
uniform token mjc:option:jacobian = "auto" (
allowedTokens = ["auto", "dense", "sparse"]
@@ -265,7 +270,9 @@ class "MjcSceneAPI" (
)
uniform double mjc:option:timestep = 0.002 (
displayName = "Timestep"
doc = "Controls the timestep in seconds used by MuJoCo."
doc = """DEPRECATED: Use newton:timeStepsPerSecond instead.
Controls the timestep in seconds used by MuJoCo."""
)
uniform double mjc:option:tolerance = 1e-8 (
displayName = "Solver Tolerance"
@@ -303,6 +310,7 @@ class "MjcImageableAPI" (
}
class "MjcCollisionAPI" (
apiSchemas = ["NewtonCollisionAPI"]
doc = "API describing a MuJoCo collider."
)
{
@@ -312,7 +320,9 @@ class "MjcCollisionAPI" (
)
uniform double mjc:gap = 0 (
displayName = "Gap"
doc = "Additional contact detection buffer beyond margin. Contacts are detected at distance margin + gap but forces are only generated at distance margin."
doc = """DEPRECATED: Use newton:contactGap instead.
Additional contact detection buffer beyond margin. Contacts are detected at distance margin + gap but forces are only generated at distance margin."""
)
uniform int mjc:group = 0 (
displayName = "Group"
@@ -320,7 +330,9 @@ class "MjcCollisionAPI" (
)
uniform double mjc:margin = 0 (
displayName = "Margin"
doc = "Geometric inflation of the geom surface for the purpose of contact force generation."
doc = """DEPRECATED: Use newton:contactMargin and newton:contactGap instead.
Geometric inflation of the geom surface for the purpose of contact force generation."""
)
uniform int mjc:priority = 0 (
displayName = "Priority"
@@ -345,6 +357,7 @@ class "MjcCollisionAPI" (
}
class "MjcMeshCollisionAPI" (
apiSchemas = ["NewtonMeshCollisionAPI"]
doc = "API describing a MuJoCo mesh collider."
)
{
@@ -355,7 +368,9 @@ class "MjcMeshCollisionAPI" (
)
uniform int mjc:maxhullvert = -1 (
displayName = "Maximum Hull Vertices"
doc = "Sets an upper limit on the number of vertices in the meshes convex hull. The default value of -1 means unlimited."
doc = """DEPRECATED: Use newton:maxHullVertices instead.
Sets an upper limit on the number of vertices in the meshes convex hull. The default value of -1 means unlimited."""
)
}
@@ -537,16 +552,23 @@ class "MjcJointAPI" (
}
class "MjcMaterialAPI" (
doc = "API providing extension attributes to represent physical MuJoCo materials."
apiSchemas = ["NewtonMaterialAPI"]
doc = """DEPRECATED: Use NewtonMaterialAPI instead. All attributes on this API have been superseded by Newton equivalents.
API providing extension attributes to represent physical MuJoCo materials."""
)
{
uniform double mjc:rollingfriction = 0.0001 (
displayName = "Rolling Friction"
doc = "Friction value acting around both axes on the contact tangent plane."
doc = """DEPRECATED: Use newton:rollingFriction instead.
Friction value acting around both axes on the contact tangent plane."""
)
uniform double mjc:torsionalfriction = 0.005 (
displayName = "Torsional Friction"
doc = "Friction value acting around contact normal."
doc = """DEPRECATED: Use newton:torsionalFriction instead.
Friction value acting around contact normal."""
)
}
@@ -586,23 +608,36 @@ class "MjcEqualityWeldAPI" (
}
class "MjcEqualityJointAPI" (
apiSchemas = ["MjcEqualityAPI"]
apiSchemas = ["NewtonMimicAPI"]
doc = """API providing extension attributes to represent equality/joint constraints.
This API is applied to a joint prim which acts as the constrained joint (joint1 in
MuJoCo terminology). The target relationship points to another joint prim which is
the reference joint (joint2 in MuJoCo terminology). The constrained joint's position
or angle is constrained to be a quartic polynomial of the reference joint's position
or angle. Only scalar joint types (slide and hinge) can be used."""
This API is applied to a joint prim which acts as the follower (joint0). The leader
joint (joint1) is specified via the newton:mimicJoint relationship inherited from
NewtonMimicAPI.
The follower's position or angle is constrained to be a quartic polynomial of the
leader's position or angle:
joint0 = coef0 + coef1*(joint1) + coef2*(joint1)^2 + coef3*(joint1)^3 + coef4*(joint1)^4
The constant (coef0) and linear (coef1) coefficients are provided by NewtonMimicAPI
as newton:mimicCoef0 and newton:mimicCoef1. The higher-order coefficients (coef2-coef4)
are provided by this API.
Only scalar joint types (slide and hinge) can be used."""
)
{
uniform double mjc:coef0 = 0 (
displayName = "Coefficient 0"
doc = """Constant coefficient a0 of the quartic polynomial. The constraint is:
doc = """DEPRECATED: Use newton:mimicCoef0 instead.
Constant coefficient a0 of the quartic polynomial. The constraint is:
y = y0 + a0 + a1*(x-x0) + a2*(x-x0)^2 + a3*(x-x0)^3 + a4*(x-x0)^4."""
)
uniform double mjc:coef1 = 1 (
displayName = "Coefficient 1"
doc = "Linear coefficient a1 of the quartic polynomial."
doc = """DEPRECATED: Use newton:mimicCoef1 instead.
Linear coefficient a1 of the quartic polynomial."""
)
uniform double mjc:coef2 = 0 (
displayName = "Coefficient 2"
@@ -616,6 +651,19 @@ class "MjcEqualityJointAPI" (
displayName = "Coefficient 4"
doc = "Quartic coefficient a4 of the quartic polynomial."
)
uniform double[] mjc:solimp = [0.9, 0.95, 0.001, 0.5, 2] (
displayName = "SolImp"
doc = "Constraint solver parameter for equality constraint simulation."
)
uniform double[] mjc:solref = [0.02, 1] (
displayName = "SolRef"
doc = "Constraint solver parameter for equality constraint simulation."
)
rel mjc:target (
doc = """DEPRECATED: Use newton:mimicJoint instead.
Secondary target of the equality constraint (the leader/reference joint)."""
)
}
class MjcTendon "MjcTendon" (
src/experimental/usd/mjcPhysics/schema.usda
+77 -17
View File
@@ -100,6 +100,7 @@ class "MjcSceneAPI"
}
doc = """API providing global simulation options for MuJoCo."""
prepend apiSchemas = ["NewtonSceneAPI"]
inherits = </APISchemaBase>
)
{
@@ -108,7 +109,9 @@ class "MjcSceneAPI"
string apiName = "Timestep"
}
displayName = "Timestep"
doc = """Controls the timestep in seconds used by MuJoCo."""
doc = """DEPRECATED: Use newton:timeStepsPerSecond instead.
Controls the timestep in seconds used by MuJoCo."""
)
uniform double mjc:option:impratio = 1.0 (
@@ -229,7 +232,9 @@ class "MjcSceneAPI"
string apiName = "Iterations"
}
displayName = "Solver Iterations"
doc = """Maximum number of iterations of the constraint solver."""
doc = """DEPRECATED: Use newton:maxSolverIterations instead.
Maximum number of iterations of the constraint solver."""
)
uniform double mjc:option:tolerance = 1e-08 (
@@ -378,7 +383,9 @@ class "MjcSceneAPI"
string apiName = "GravityFlag"
}
displayName = "Gravity Toggle"
doc = """Enables the application of gravitational acceleration as defined in mjOption."""
doc = """DEPRECATED: Use newton:gravityEnabled instead.
Enables the application of gravitational acceleration as defined in mjOption."""
)
uniform bool mjc:flag:clampctrl = True (
@@ -674,6 +681,7 @@ class "MjcCollisionAPI"
}
doc = """API describing a MuJoCo collider."""
prepend apiSchemas = ["NewtonCollisionAPI"]
inherits = </APISchemaBase>
)
{
@@ -738,7 +746,9 @@ class "MjcCollisionAPI"
string apiName = "Margin"
}
displayName = "Margin"
doc = """Geometric inflation of the geom surface for the purpose of contact force generation."""
doc = """DEPRECATED: Use newton:contactMargin and newton:contactGap instead.
Geometric inflation of the geom surface for the purpose of contact force generation."""
)
uniform double mjc:gap = 0.0 (
@@ -746,7 +756,9 @@ class "MjcCollisionAPI"
string apiName = "Gap"
}
displayName = "Gap"
doc = """Additional contact detection buffer beyond margin. Contacts are detected at distance margin + gap but forces are only generated at distance margin."""
doc = """DEPRECATED: Use newton:contactGap instead.
Additional contact detection buffer beyond margin. Contacts are detected at distance margin + gap but forces are only generated at distance margin."""
)
}
@@ -757,6 +769,7 @@ class "MjcMeshCollisionAPI"
}
doc = """API describing a MuJoCo mesh collider."""
prepend apiSchemas = ["NewtonMeshCollisionAPI"]
inherits = </APISchemaBase>
)
{
@@ -774,7 +787,9 @@ class "MjcMeshCollisionAPI"
string apiName = "MaxHullVert"
}
displayName = "Maximum Hull Vertices"
doc = """Sets an upper limit on the number of vertices in the meshes convex hull. The default value of -1 means unlimited."""
doc = """DEPRECATED: Use newton:maxHullVertices instead.
Sets an upper limit on the number of vertices in the meshes convex hull. The default value of -1 means unlimited."""
)
}
@@ -1035,8 +1050,11 @@ class "MjcMaterialAPI"
customData = {
string className = "MaterialAPI"
}
doc = """API providing extension attributes to represent physical MuJoCo materials."""
doc = """DEPRECATED: Use NewtonMaterialAPI instead. All attributes on this API have been superseded by Newton equivalents.
API providing extension attributes to represent physical MuJoCo materials."""
prepend apiSchemas = ["NewtonMaterialAPI"]
inherits = </APISchemaBase>
)
{
@@ -1045,7 +1063,9 @@ class "MjcMaterialAPI"
string apiName = "TorsionalFriction"
}
displayName = "Torsional Friction"
doc = """Friction value acting around contact normal."""
doc = """DEPRECATED: Use newton:torsionalFriction instead.
Friction value acting around contact normal."""
)
uniform double mjc:rollingfriction = 0.0001 (
@@ -1053,7 +1073,9 @@ class "MjcMaterialAPI"
string apiName = "RollingFriction"
}
displayName = "Rolling Friction"
doc = """Friction value acting around both axes on the contact tangent plane."""
doc = """DEPRECATED: Use newton:rollingFriction instead.
Friction value acting around both axes on the contact tangent plane."""
)
}
@@ -1121,22 +1143,58 @@ class "MjcEqualityJointAPI" (
string className = "EqualityJointAPI"
}
doc = """API providing extension attributes to represent equality/joint constraints.
This API is applied to a joint prim which acts as the constrained joint (joint1 in
MuJoCo terminology). The target relationship points to another joint prim which is
the reference joint (joint2 in MuJoCo terminology). The constrained joint's position
or angle is constrained to be a quartic polynomial of the reference joint's position
or angle. Only scalar joint types (slide and hinge) can be used."""
prepend apiSchemas = ["MjcEqualityAPI"]
This API is applied to a joint prim which acts as the follower (joint0). The leader
joint (joint1) is specified via the newton:mimicJoint relationship inherited from
NewtonMimicAPI.
The follower's position or angle is constrained to be a quartic polynomial of the
leader's position or angle:
joint0 = coef0 + coef1*(joint1) + coef2*(joint1)^2 + coef3*(joint1)^3 + coef4*(joint1)^4
The constant (coef0) and linear (coef1) coefficients are provided by NewtonMimicAPI
as newton:mimicCoef0 and newton:mimicCoef1. The higher-order coefficients (coef2-coef4)
are provided by this API.
Only scalar joint types (slide and hinge) can be used."""
prepend apiSchemas = ["NewtonMimicAPI"]
inherits = </APISchemaBase>
)
{
uniform double[] mjc:solimp = [0.9, 0.95, 0.001, 0.5, 2] (
customData = {
string apiName = "SolImp"
}
displayName = "SolImp"
doc = """Constraint solver parameter for equality constraint simulation."""
)
uniform double[] mjc:solref = [0.02, 1] (
customData = {
string apiName = "SolRef"
}
displayName = "SolRef"
doc = """Constraint solver parameter for equality constraint simulation."""
)
rel mjc:target (
customData = {
string apiName = "MjcTarget"
}
doc = """DEPRECATED: Use newton:mimicJoint instead.
Secondary target of the equality constraint (the leader/reference joint)."""
)
uniform double mjc:coef0 = 0 (
customData = {
string apiName = "Coef0"
}
displayName = "Coefficient 0"
doc = """Constant coefficient a0 of the quartic polynomial. The constraint is:
doc = """DEPRECATED: Use newton:mimicCoef0 instead.
Constant coefficient a0 of the quartic polynomial. The constraint is:
y = y0 + a0 + a1*(x-x0) + a2*(x-x0)^2 + a3*(x-x0)^3 + a4*(x-x0)^4."""
)
@@ -1145,7 +1203,9 @@ class "MjcEqualityJointAPI" (
string apiName = "Coef1"
}
displayName = "Coefficient 1"
doc = """Linear coefficient a1 of the quartic polynomial."""
doc = """DEPRECATED: Use newton:mimicCoef1 instead.
Linear coefficient a1 of the quartic polynomial."""
)
uniform double mjc:coef2 = 0 (
src/experimental/usd/plugins/mjcf/mujoco_to_usd.cc
+46 -20
View File
@@ -15,6 +15,7 @@
#include "mjcf/mujoco_to_usd.h"
#include <algorithm>
#include <cmath>
#include <cstddef>
#include <numbers>
#include <string>
@@ -119,6 +120,16 @@ TF_DEFINE_PRIVATE_TOKENS(kTokens,
(UsdPrimvarReader_float2)
(UsdUVTexture)
(UsdPreviewSurface)
((NewtonMaterialAPI, "NewtonMaterialAPI"))
((NewtonMeshCollisionAPI, "NewtonMeshCollisionAPI"))
((newtonTorsionalFriction, "newton:torsionalFriction"))
((newtonRollingFriction, "newton:rollingFriction"))
((newtonMaxHullVertices, "newton:maxHullVertices"))
((newtonMaxSolverIterations, "newton:maxSolverIterations"))
((newtonTimeStepsPerSecond, "newton:timeStepsPerSecond"))
((newtonGravityEnabled, "newton:gravityEnabled"))
((newtonContactMargin, "newton:contactMargin"))
((newtonContactGap, "newton:contactGap"))
);
// Using to satisfy TF_REGISTRY_FUNCTION macro below and avoid operating in PXR_NS.
@@ -393,8 +404,11 @@ class ModelWriter {
WriteUniformAttribute(mesh_spec, pxr::SdfValueTypeNames->Token,
MjcPhysicsTokens->mjcInertia, inertia);
WriteUniformAttribute(mesh_spec, pxr::SdfValueTypeNames->Int,
MjcPhysicsTokens->mjcMaxhullvert, mesh->maxhullvert);
// Newton mesh attribute (replaces deprecated mjc:maxhullvert)
if (mesh->maxhullvert != -1) {
WriteUniformAttribute(mesh_spec, pxr::SdfValueTypeNames->Int,
kTokens->newtonMaxHullVertices, mesh->maxhullvert);
}
// NOTE: The geometry data taken from the spec is the post-compilation
// data after it has been mjCMesh::Compile'd. So don't be surprised if
@@ -500,7 +514,7 @@ class ModelWriter {
const std::vector<std::pair<pxr::TfToken, double>>
option_double_attributes = {
{MjcPhysicsTokens->mjcOptionTimestep, spec_->option.timestep},
// mjc:option:timestep deprecated in favor of newton:timeStepsPerSecond
{MjcPhysicsTokens->mjcOptionTolerance, spec_->option.tolerance},
{MjcPhysicsTokens->mjcOptionLs_tolerance,
spec_->option.ls_tolerance},
@@ -519,7 +533,7 @@ class ModelWriter {
}
const std::vector<std::pair<pxr::TfToken, int>> option_int_attributes = {
{MjcPhysicsTokens->mjcOptionIterations, spec_->option.iterations},
// mjc:option:iterations deprecated in favor of newton:maxSolverIterations
{MjcPhysicsTokens->mjcOptionLs_iterations, spec_->option.ls_iterations},
{MjcPhysicsTokens->mjcOptionNoslip_iterations,
spec_->option.noslip_iterations},
@@ -665,7 +679,7 @@ class ModelWriter {
{MjcPhysicsTokens->mjcFlagContact, mjDSBL_CONTACT},
{MjcPhysicsTokens->mjcFlagSpring, mjDSBL_SPRING},
{MjcPhysicsTokens->mjcFlagDamper, mjDSBL_DAMPER},
{MjcPhysicsTokens->mjcFlagGravity, mjDSBL_GRAVITY},
// mjc:flag:gravity deprecated in favor of newton:gravityEnabled
{MjcPhysicsTokens->mjcFlagClampctrl, mjDSBL_CLAMPCTRL},
{MjcPhysicsTokens->mjcFlagWarmstart, mjDSBL_WARMSTART},
{MjcPhysicsTokens->mjcFlagFilterparent, mjDSBL_FILTERPARENT},
@@ -747,6 +761,19 @@ class ModelWriter {
WriteUniformAttribute(physics_scene_spec, pxr::SdfValueTypeNames->Bool,
MjcPhysicsTokens->mjcCompilerSaveInertial,
(bool)spec_->compiler.saveinertial);
// Newton scene attributes (auto-applied via MjcSceneAPI -> NewtonSceneAPI)
WriteUniformAttribute(physics_scene_spec, pxr::SdfValueTypeNames->Int,
kTokens->newtonMaxSolverIterations,
spec_->option.iterations);
if (spec_->option.timestep > 0) {
WriteUniformAttribute(physics_scene_spec, pxr::SdfValueTypeNames->Int,
kTokens->newtonTimeStepsPerSecond,
static_cast<int>(std::round(1.0 / spec_->option.timestep)));
}
bool gravity_disabled = spec_->option.disableflags & mjDSBL_GRAVITY;
WriteUniformAttribute(physics_scene_spec, pxr::SdfValueTypeNames->Bool,
kTokens->newtonGravityEnabled, !gravity_disabled);
}
void WriteMeshes() {
@@ -868,16 +895,13 @@ class ModelWriter {
pxr::UsdPhysicsTokens->physicsDynamicFriction,
(float)geom->friction[0]);
}
if (geom->friction[1] != geom_default->friction[1]) {
WriteUniformAttribute(material_spec, pxr::SdfValueTypeNames->Double,
MjcPhysicsTokens->mjcTorsionalfriction,
geom->friction[1]);
}
if (geom->friction[2] != geom_default->friction[2]) {
WriteUniformAttribute(material_spec, pxr::SdfValueTypeNames->Double,
MjcPhysicsTokens->mjcRollingfriction,
geom->friction[2]);
}
// Newton material attributes (replaces deprecated mjc:torsionalfriction / mjc:rollingfriction)
WriteUniformAttribute(material_spec, pxr::SdfValueTypeNames->Float,
kTokens->newtonTorsionalFriction,
(float)geom->friction[1]);
WriteUniformAttribute(material_spec, pxr::SdfValueTypeNames->Float,
kTokens->newtonRollingFriction,
(float)geom->friction[2]);
return material_spec;
}
@@ -1706,11 +1730,13 @@ class ModelWriter {
MjcPhysicsTokens->mjcSolimp,
pxr::VtArray<double>(geom->solimp, geom->solimp + mjNIMP));
WriteUniformAttribute(geom_spec, pxr::SdfValueTypeNames->Double,
MjcPhysicsTokens->mjcMargin, geom->margin);
WriteUniformAttribute(geom_spec, pxr::SdfValueTypeNames->Double,
MjcPhysicsTokens->mjcGap, geom->gap);
// Newton collision attributes (replaces deprecated mjc:margin / mjc:gap)
WriteUniformAttribute(geom_spec, pxr::SdfValueTypeNames->Float,
kTokens->newtonContactMargin,
static_cast<float>(geom->margin));
WriteUniformAttribute(geom_spec, pxr::SdfValueTypeNames->Float,
kTokens->newtonContactGap,
static_cast<float>(geom->gap));
if (geom->mass >= mjMINVAL || geom->density >= mjMINVAL) {
ApplyApiSchema(layer_, geom_spec, pxr::UsdPhysicsTokens->PhysicsMassAPI);
test/engine/engine_collision_gjk_test.cc
+338 -484
View File
File diff suppressed because it is too large Load Diff
test/engine/engine_core_smooth_test.cc
+19 -4
View File
@@ -233,13 +233,13 @@ TEST_F(CoreSmoothTest, TendonArmature) {
// get full M, includes both CRB and tendon inertia
vector<mjtNum> M(nv*nv);
mj_fullM(m, M.data(), d->qM);
mj_fullM(m, d, M.data());
// put only CRB inertia in M2
mj_crb(m, d);
mju_scatter(d->qM, d->M, m->mapM2M, m->nC);
vector<mjtNum> M2(nv*nv);
mj_fullM(m, M2.data(), d->qM);
mj_fullM(m, d, M2.data());
vector<mjtNum> ten_J(nv); // tendon Jacobian
vector<mjtNum> ten_M(nv*nv); // tendon inertia
@@ -681,7 +681,7 @@ TEST_F(CoreSmoothTest, FactorI) {
// dense M matrix
vector<mjtNum> Mexpected(nv*nv);
mj_fullM(model, Mexpected.data(), data->qM);
mj_fullM(model, data, Mexpected.data());
// expect matrices to match to floating point precision
EXPECT_THAT(M, Pointwise(MjNear(1e-12, 1e-5), Mexpected));
@@ -690,6 +690,21 @@ TEST_F(CoreSmoothTest, FactorI) {
mj_deleteModel(model);
}
// Convert legacy-format symmetric matrix to dense (local helper for tests).
static void legacyToDense(const mjModel* m, mjtNum* dst, const mjtNum* M) {
int adr = 0, nv = m->nv;
mju_zero(dst, nv*nv);
for (int i = 0; i < nv; i++) {
int j = i;
while (j >= 0) {
dst[i*nv+j] = M[adr];
dst[j*nv+i] = M[adr];
j = m->dof_parentid[j];
adr++;
}
}
}
TEST_F(CoreSmoothTest, SolveLDs) {
const std::string xml_path = GetTestDataFilePath(kInertiaPath);
char error[1024];
@@ -712,7 +727,7 @@ TEST_F(CoreSmoothTest, SolveLDs) {
mju_sparse2dense(LDdense.data(), d->qLD, nv, nv,
m->M_rownnz, m->M_rowadr, m->M_colind);
vector<mjtNum> LDdense2(nv*nv);
mj_fullM(m, LDdense2.data(), LDlegacy.data());
legacyToDense(m, LDdense2.data(), LDlegacy.data());
// expect lower triangles to match exactly
for (int i=0; i < nv; i++) {
test/engine/engine_derivative_test.cc
+1 -1
View File
@@ -848,7 +848,7 @@ TEST_F(DerivativeTest, LinearSystemInverse) {
// expect that acceleration derivatives are the mass matrix
vector<mjtNum> DfDa_expect(nv*nv, 0);
mj_fullM(model, DfDa_expect.data(), data->qM);
mj_fullM(model, data, DfDa_expect.data());
EXPECT_THAT(DfDa, Pointwise(DoubleNear(eps), DfDa_expect));
// expect that sensor derivatives w.r.t position only see sensor 1 at dof 0
test/engine/engine_support_test.cc
+3 -3
View File
@@ -608,13 +608,13 @@ TEST_F(InertiaTest, FullM) {
ASSERT_THAT(m, NotNull()) << "Failed to load model: " << error;
int nv = m->nv;
// forward dynamics, populate qM and qLD
// forward dynamics, populate M and qLD
mjData* d = mj_makeData(m);
mj_forward(m, d);
// get dense mass matrix from M using mju_sym2dense
// get dense mass matrix from M using mj_fullM
vector<mjtNum> M(nv * nv);
mju_sym2dense(M.data(), d->M, nv, m->M_rownnz, m->M_rowadr, m->M_colind);
mj_fullM(m, d, M.data());
// get dense mass matrix from M using mju_sparse2dense
vector<mjtNum> M_CSR(nv * nv);
test/engine/engine_util_errmem_test.cc
-64
View File
@@ -68,70 +68,6 @@ class MujocoErrorAndWarningTest : public ::testing::Test {
}
};
TEST_F(MujocoErrorAndWarningTest, MjuErrorI) {
std::string format_string = "%010d";
while (format_string.length() < 2 * kBufferSize) {
format_string += 'x';
}
std::string expected_message = "0123456789";
while (expected_message.length() < kBufferSize - 1) {
expected_message += 'x';
}
ClearErrorMessage();
mju_error_i(format_string.c_str(), 123456789);
EXPECT_EQ(std::string(ErrorMessageBuffer()), expected_message);
}
TEST_F(MujocoErrorAndWarningTest, MjuWarningI) {
std::string format_string = "%010d";
while (format_string.length() < 2 * kBufferSize) {
format_string += 'x';
}
std::string expected_message = "0123456789";
while (expected_message.length() < kBufferSize - 1) {
expected_message += 'x';
}
ClearWarningMessage();
mju_warning_i(format_string.c_str(), 123456789);
EXPECT_EQ(std::string(WarningMessageBuffer()), expected_message);
}
TEST_F(MujocoErrorAndWarningTest, MjuErrorS) {
std::string format_string = "% 9s";
while (format_string.length() < 2 * kBufferSize) {
format_string += 'z';
}
std::string expected_message = " foobar";
while (expected_message.length() < kBufferSize - 1) {
expected_message += 'z';
}
ClearErrorMessage();
mju_error_s(format_string.c_str(), "foobar");
EXPECT_EQ(std::string(ErrorMessageBuffer()), expected_message);
}
TEST_F(MujocoErrorAndWarningTest, MjuWarningS) {
std::string format_string = "% 9s";
while (format_string.length() < 2 * kBufferSize) {
format_string += 'z';
}
std::string expected_message = " foobar";
while (expected_message.length() < kBufferSize - 1) {
expected_message += 'z';
}
ClearWarningMessage();
mju_warning_s(format_string.c_str(), "foobar");
EXPECT_EQ(std::string(WarningMessageBuffer()), expected_message);
}
TEST_F(MujocoErrorAndWarningTest, MjuErrorInternal) {
ClearErrorMessage();
mjERROR("foobar %d", 123);
test/experimental/usd/plugins/mjcf/mjcf_file_format_test.cc
+42 -12
View File
@@ -78,6 +78,9 @@ PXR_NAMESPACE_OPEN_SCOPE
// clang-format off
TF_DEFINE_PRIVATE_TOKENS(_tokens,
(st)
((newtonTimeStepsPerSecond, "newton:timeStepsPerSecond"))
((newtonMaxSolverIterations, "newton:maxSolverIterations"))
((newtonGravityEnabled, "newton:gravityEnabled"))
);
// clang-format on
PXR_NAMESPACE_CLOSE_SCOPE
@@ -159,12 +162,20 @@ TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsMaterials) {
4.0f);
ExpectAttributeEqual(stage,
"/physics_materials_test/PhysicsMaterials/"
"geom_with_friction.mjc:torsionalfriction",
5.0);
"geom_with_friction.newton:torsionalFriction",
5.0f);
ExpectAttributeEqual(stage,
"/physics_materials_test/PhysicsMaterials/"
"geom_with_friction.mjc:rollingfriction",
6.0);
"geom_with_friction.newton:rollingFriction",
6.0f);
EXPECT_ATTRIBUTE_HAS_NO_AUTHORED_VALUE(
stage,
"/physics_materials_test/PhysicsMaterials/"
"geom_with_friction.mjc:torsionalfriction");
EXPECT_ATTRIBUTE_HAS_NO_AUTHORED_VALUE(
stage,
"/physics_materials_test/PhysicsMaterials/"
"geom_with_friction.mjc:rollingfriction");
}
TEST_F(MjcfSdfFileFormatPluginTest, TestMaterials) {
@@ -692,10 +703,15 @@ TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsScenePrimTimestep) {
</mujoco>
)");
// newton:timeStepsPerSecond = round(1/0.005) = 200
ExpectAttributeEqual(
stage,
kPhysicsScenePrimPath.AppendProperty(MjcPhysicsTokens->mjcOptionTimestep),
0.005);
kPhysicsScenePrimPath.AppendProperty(pxr::_tokens->newtonTimeStepsPerSecond),
200);
// deprecated mjc:option:timestep should not be authored
EXPECT_ATTRIBUTE_HAS_NO_AUTHORED_VALUE(
stage,
kPhysicsScenePrimPath.AppendProperty(MjcPhysicsTokens->mjcOptionTimestep));
}
TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsScenePrimCone) {
@@ -928,8 +944,12 @@ TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsScenePrimIterations) {
ExpectAttributeEqual(stage,
kPhysicsScenePrimPath.AppendProperty(
MjcPhysicsTokens->mjcOptionIterations),
pxr::_tokens->newtonMaxSolverIterations),
10);
EXPECT_ATTRIBUTE_HAS_NO_AUTHORED_VALUE(
stage,
kPhysicsScenePrimPath.AppendProperty(
MjcPhysicsTokens->mjcOptionIterations));
}
TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsScenePrimLSIterations) {
@@ -1066,7 +1086,7 @@ TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsScenePrimDisableFlags) {
MjcPhysicsTokens->mjcFlagContact,
MjcPhysicsTokens->mjcFlagSpring,
MjcPhysicsTokens->mjcFlagDamper,
MjcPhysicsTokens->mjcFlagGravity,
// mjc:flag:gravity is deprecated, now exported as newton:gravityEnabled
MjcPhysicsTokens->mjcFlagClampctrl,
MjcPhysicsTokens->mjcFlagWarmstart,
MjcPhysicsTokens->mjcFlagFilterparent,
@@ -1083,6 +1103,10 @@ TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsScenePrimDisableFlags) {
ExpectAttributeEqual(stage, kPhysicsScenePrimPath.AppendProperty(flag),
false);
}
ExpectAttributeEqual(
stage,
kPhysicsScenePrimPath.AppendProperty(pxr::_tokens->newtonGravityEnabled),
false);
}
TEST_F(MjcfSdfFileFormatPluginTest, TestPhysicsScenePrimEnableFlags) {
@@ -1472,8 +1496,11 @@ TEST_F(MjcfSdfFileFormatPluginTest, TestMjcPhysicsCollisionAPI) {
pxr::VtArray<double>({0.1, 0.2}));
ExpectAttributeEqual(stage, "/test/body/box.mjc:solimp",
pxr::VtArray<double>({0.3, 0.4, 0.5, 0.6, 0.7}));
ExpectAttributeEqual(stage, "/test/body/box.mjc:margin", 0.8);
ExpectAttributeEqual(stage, "/test/body/box.mjc:gap", 0.9);
// margin and gap are now exported as Newton attributes
EXPECT_ATTRIBUTE_HAS_NO_AUTHORED_VALUE(stage, "/test/body/box.mjc:margin");
EXPECT_ATTRIBUTE_HAS_NO_AUTHORED_VALUE(stage, "/test/body/box.mjc:gap");
ExpectAttributeEqual(stage, "/test/body/box.newton:contactMargin", 0.8f);
ExpectAttributeEqual(stage, "/test/body/box.newton:contactGap", 0.9f);
ExpectAttributeEqual(stage, "/test/body/box.mjc:shellinertia", true);
}
@@ -1508,8 +1535,11 @@ TEST_F(MjcfSdfFileFormatPluginTest, TestMjcPhysicsMeshCollisionAPI) {
MjcPhysicsTokens->convex);
ExpectAttributeEqual(stage, "/test/body/tet_shell/Mesh.mjc:inertia",
MjcPhysicsTokens->shell);
ExpectAttributeEqual(stage, "/test/body/tet_max_vert/Mesh.mjc:maxhullvert",
12);
// mjc:maxhullvert deprecated, newton:maxHullVertices used instead
EXPECT_ATTRIBUTE_HAS_NO_AUTHORED_VALUE(
stage, "/test/body/tet_max_vert/Mesh.mjc:maxhullvert");
ExpectAttributeEqual(
stage, "/test/body/tet_max_vert/Mesh.newton:maxHullVertices", 12);
}
TEST_F(MjcfSdfFileFormatPluginTest, TestMassAPIApplied) {
unity/Runtime/Bindings/MjBindings.cs
+1 -13
View File
@@ -6946,7 +6946,7 @@ public static unsafe extern int mj_name2id(mjModel_* m, int type, [MarshalAs(Unm
public static unsafe extern IntPtr mj_id2name(mjModel_* m, int type, int id);
[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
public static unsafe extern void mj_fullM(mjModel_* m, double* dst, double* M);
public static unsafe extern void mj_fullM(mjModel_* m, mjData_* d, double* dst);
[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
public static unsafe extern void mj_mulM(mjModel_* m, mjData_* d, double* res, double* vec);
@@ -7224,21 +7224,9 @@ public static unsafe extern void mjui_render(mjUI_* ui, mjuiState_* state, mjrCo
[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
public static unsafe extern void mju_error([MarshalAs(UnmanagedType.LPStr)]string msg);
[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
public static unsafe extern void mju_error_i([MarshalAs(UnmanagedType.LPStr)]string msg, int i);
[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
public static unsafe extern void mju_error_s([MarshalAs(UnmanagedType.LPStr)]string msg, [MarshalAs(UnmanagedType.LPStr)]string text);
[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
public static unsafe extern void mju_warning([MarshalAs(UnmanagedType.LPStr)]string msg);
[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
public static unsafe extern void mju_warning_i([MarshalAs(UnmanagedType.LPStr)]string msg, int i);
[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
public static unsafe extern void mju_warning_s([MarshalAs(UnmanagedType.LPStr)]string msg, [MarshalAs(UnmanagedType.LPStr)]string text);
[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
public static unsafe extern void mju_clearHandlers();
wasm/codegen/generated/bindings.cc
+2 -4
View File
@@ -8709,12 +8709,10 @@ void mj_forwardSkip_wrapper(const MjModel& m, MjData& d, int skipstage, int skip
mj_forwardSkip(m.get(), d.get(), skipstage, skipsensor);
}
void mj_fullM_wrapper(const MjModel& m, const val& dst, const NumberArray& M) {
void mj_fullM_wrapper(const MjModel& m, const MjData& d, const val& dst) {
UNPACK_VALUE(mjtNum, dst);
UNPACK_ARRAY(mjtNum, M);
CHECK_SIZE(M, m.nM());
CHECK_SIZE(dst, m.nv() * m.nv());
mj_fullM(m.get(), dst_.data(), M_.data());
mj_fullM(m.get(), d.get(), dst_.data());
}
void mj_fwdAcceleration_wrapper(const MjModel& m, MjData& d) {
wasm/codegen/generators/constants.py
-5
View File
@@ -147,14 +147,10 @@ _SKIPPED_MEMORY_FUNCTIONS: tuple[str, ...] = (
"mju_boxQPmalloc",
"mju_clearHandlers",
"mju_error",
"mju_error_i",
"mju_error_s",
"mju_free",
"mju_malloc",
"mju_strncpy",
"mju_warning",
"mju_warning_i",
"mju_warning_s",
# go/keep-sorted end
)
@@ -569,7 +565,6 @@ FUNCTION_BOUNDS_CHECKS: Dict[str, str] = {
CHECK_SIZE(qpos2, m.nq());
""".strip(),
"mj_fullM": """
CHECK_SIZE(M, m.nM());
CHECK_SIZE(dst, m.nv() * m.nv());
""".strip(),
"mj_geomDistance": """