mujodico 3.15.0__tar.gz

mujodico-3.15.0/MANIFEST.in

@@ -0,0 +1 @@
+recursive-include mujodico/mjx/test_data *

mujodico-3.15.0/PKG-INFO

@@ -0,0 +1,116 @@
+Metadata-Version: 2.4
+Name: mujodico
+Version: 3.15.0
+Summary: MuJoDiCo
+Author-email: Wei-Chen Li <williamrob104@gmail.com>
+License-Expression: Apache-2.0
+Project-URL: Homepage, https://github.com/wei-chen-li/mujodico
+Project-URL: Repository, https://github.com/wei-chen-li/mujodico
+Classifier: Development Status :: 5 - Production/Stable
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: Natural Language :: English
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Classifier: Programming Language :: Python :: 3.15
+Classifier: Topic :: Scientific/Engineering
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: absl-py
+Requires-Dist: etils[epath]
+Requires-Dist: jax
+Requires-Dist: jaxlib
+Requires-Dist: mujoco>=3.15.0.dev0
+Requires-Dist: scipy
+Requires-Dist: trimesh
+Provides-Extra: warp
+Requires-Dist: warp-lang==1.17.0; extra == "warp"
+Provides-Extra: dev
+Requires-Dist: isort; extra == "dev"
+Requires-Dist: pyink; extra == "dev"
+
+# MuJoDiCo (Multi-Joint dynamics with Differentiable Contact)
+
+This package extends MuJoCo with differentiable contact dynamics and provides trajectory optimization algorithms for contact-rich dynamical systems. Gradients are computed using the implicit function theorem rather than differentiating through the solver rollout.
+
+## Installation
+
+The recommended way to install this package is via [PyPI](https://pypi.org/project/mujodico/):
+
+```sh
+pip install mujodico
+```
+
+For local development, install from a checkout:
+
+```sh
+git clone https://github.com/wei-chen-li/mujodico.git
+cd mujodico
+pip install -e .
+```
+
+## Usage
+
+[MuJoDiCo](https://github.com/wei-chen-li/mujodico) is a drop-in replacement for [MuJoCo MJX](https://mujoco.readthedocs.io/en/stable/mjx.html) that provides differentiable contact dynamics. You can simply replace `from mujoco import mjx` with `from mujodico import mjx`. For example:
+
+```python
+import jax
+import jax.numpy as jnp
+import mujoco
+from mujodico import mjx
+
+model = mujoco.MjModel.from_xml_string(XML)
+
+mjx_model = mjx.put_model(
+    model,
+    contact_approximation='Hunt-Crossley',
+    contact_smoothing=0.0,
+)
+
+def f(x, u):
+    nq, nv = mjx_model.nq, mjx_model.nv
+    d = mjx.make_data(mjx_model)
+    d = d.replace(qpos=x[:nq], qvel=x[nq:nq+nv], act=x[nq+nv:], ctrl=u)
+    d = mjx.step(mjx_model, d)
+    x_next = jnp.concatenate((d.qpos, d.qvel, d.act))
+    return x_next
+
+x = jnp.zeros(mjx_model.nq + mjx_model.nv + mjx_model.na)
+u = jnp.zeros(mjx_model.nu)
+
+A, B = jax.jacfwd(f, argnums=(0,1))(x, u)
+```
+
+MuJoDiCo computes dynamics Jacobian using the implicit function theorem. Both forward-mode (`jax.jacfwd`) and reverse-mode (`jax.jacrev`) differentiation are supported.
+
+### Contact Modeling
+
+MuJoDiCo provides configurable contact models through two parameters:
+
+* **`contact_approximation`** — specifies the contact model used for the dynamics. `'Hunt-Crossley'` is recommended for most applications.
+* **`contact_smoothing`** — adds artificial smoothing to the contact model to make the resulting dynamics **C¹-continuous**.
+
+The `contact_smoothing=κ` parameter controls the amount of regularization. It can be interpreted as applying a force of κ newtons over a distance of 1 meter.
+
+For forward simulation, where differentiability is not required, use `contact_smoothing=0.0`.
+For trajectory optimization or other gradient-based methods that require C¹-continuous dynamics, `contact_smoothing=1e-5` is a good starting point, providing smooth dynamics while introducing relatively little unphysical behavior.
+
+### Examples
+
+```sh
+python -m mujodico.examples --help
+```
+
+## Versioning
+
+The `major.minor.micro` portion of the version number matches the version of MuJoCo. Optionally, if we release updates to MuJoDiCo that target the same version of MuJoCo, a `.postN` suffix is added, for example `3.0.1.post2` represents the second update to MuJoDiCo for MuJoCo 3.0.1.
+
+## License and Disclaimer
+
+MuJoDiCo is licensed under the Apache License, Version 2.0. You may obtain a copy of the License at https://www.apache.org/licenses/LICENSE-2.0.
+
+This project includes code adapted from [MuJoCo](https://github.com/google-deepmind/mujoco), which is licensed under Apache License, Version 2.0. You may obtain a copy of the License at https://www.apache.org/licenses/LICENSE-2.0.

mujodico-3.15.0/README.md

@@ -0,0 +1,81 @@
+# MuJoDiCo (Multi-Joint dynamics with Differentiable Contact)
+
+This package extends MuJoCo with differentiable contact dynamics and provides trajectory optimization algorithms for contact-rich dynamical systems. Gradients are computed using the implicit function theorem rather than differentiating through the solver rollout.
+
+## Installation
+
+The recommended way to install this package is via [PyPI](https://pypi.org/project/mujodico/):
+
+```sh
+pip install mujodico
+```
+
+For local development, install from a checkout:
+
+```sh
+git clone https://github.com/wei-chen-li/mujodico.git
+cd mujodico
+pip install -e .
+```
+
+## Usage
+
+[MuJoDiCo](https://github.com/wei-chen-li/mujodico) is a drop-in replacement for [MuJoCo MJX](https://mujoco.readthedocs.io/en/stable/mjx.html) that provides differentiable contact dynamics. You can simply replace `from mujoco import mjx` with `from mujodico import mjx`. For example:
+
+```python
+import jax
+import jax.numpy as jnp
+import mujoco
+from mujodico import mjx
+
+model = mujoco.MjModel.from_xml_string(XML)
+
+mjx_model = mjx.put_model(
+    model,
+    contact_approximation='Hunt-Crossley',
+    contact_smoothing=0.0,
+)
+
+def f(x, u):
+    nq, nv = mjx_model.nq, mjx_model.nv
+    d = mjx.make_data(mjx_model)
+    d = d.replace(qpos=x[:nq], qvel=x[nq:nq+nv], act=x[nq+nv:], ctrl=u)
+    d = mjx.step(mjx_model, d)
+    x_next = jnp.concatenate((d.qpos, d.qvel, d.act))
+    return x_next
+
+x = jnp.zeros(mjx_model.nq + mjx_model.nv + mjx_model.na)
+u = jnp.zeros(mjx_model.nu)
+
+A, B = jax.jacfwd(f, argnums=(0,1))(x, u)
+```
+
+MuJoDiCo computes dynamics Jacobian using the implicit function theorem. Both forward-mode (`jax.jacfwd`) and reverse-mode (`jax.jacrev`) differentiation are supported.
+
+### Contact Modeling
+
+MuJoDiCo provides configurable contact models through two parameters:
+
+* **`contact_approximation`** — specifies the contact model used for the dynamics. `'Hunt-Crossley'` is recommended for most applications.
+* **`contact_smoothing`** — adds artificial smoothing to the contact model to make the resulting dynamics **C¹-continuous**.
+
+The `contact_smoothing=κ` parameter controls the amount of regularization. It can be interpreted as applying a force of κ newtons over a distance of 1 meter.
+
+For forward simulation, where differentiability is not required, use `contact_smoothing=0.0`.
+For trajectory optimization or other gradient-based methods that require C¹-continuous dynamics, `contact_smoothing=1e-5` is a good starting point, providing smooth dynamics while introducing relatively little unphysical behavior.
+
+### Examples
+
+```sh
+python -m mujodico.examples --help
+```
+
+## Versioning
+
+The `major.minor.micro` portion of the version number matches the version of MuJoCo. Optionally, if we release updates to MuJoDiCo that target the same version of MuJoCo, a `.postN` suffix is added, for example `3.0.1.post2` represents the second update to MuJoDiCo for MuJoCo 3.0.1.
+
+## License and Disclaimer
+
+MuJoDiCo is licensed under the Apache License, Version 2.0. You may obtain a copy of the License at https://www.apache.org/licenses/LICENSE-2.0.
+
+This project includes code adapted from [MuJoCo](https://github.com/google-deepmind/mujoco), which is licensed under Apache License, Version 2.0. You may obtain a copy of the License at https://www.apache.org/licenses/LICENSE-2.0.

mujodico-3.15.0/mujodico/examples/__main__.py

@@ -0,0 +1,107 @@
+import os
+import time
+import csv
+from pathlib import Path
+
+from absl import app
+from absl import flags
+import jax
+from jax import numpy as jp
+import mujoco
+from mujodico import mjx
+from mujodico.examples.methods import get_method_names, parse_method_name, support
+from mujodico.examples.tasks import get_task_names, parse_task_name
+from mujodico.examples.visualizer import Visualizer
+import numpy as np
+
+
+_METHOD = flags.DEFINE_string(
+    'method',
+    None,
+    f'<{"|".join(get_method_names())}>\n(required)',
+    required=True,
+)
+_TASK = flags.DEFINE_string(
+    'task',
+    None,
+    f'<{"|".join(get_task_names())}>\n(required)',
+    required=True,
+)
+_ITERATIONS = flags.DEFINE_integer(
+    'iterations', 500, 'Number of solver iterations', lower_bound=0,
+)
+_TESTS = flags.DEFINE_integer(
+    'tests', 1, 'Number of test cases to run.', lower_bound=1,
+)
+_SEED = flags.DEFINE_integer(
+    'seed', 0, 'Random seed for tests.',
+)
+
+
+def _main(argv) -> None:
+  if len(argv) > 1:
+    raise app.UsageError('Too many command-line arguments.')
+
+  method_cls = parse_method_name(_METHOD.value)
+  task_cls = parse_task_name(_TASK.value)
+
+  result_dir = Path(__file__).resolve().parent / 'result'
+  result_dir.mkdir(exist_ok=True)
+  result_file = result_dir / f'{task_cls.__name__}_{method_cls.__name__}.csv'
+  print(f'Writing results to "{result_file}"')
+
+  with open(result_file, 'a', newline='') as f:
+    writer = csv.writer(f)
+    if f.tell() == 0:
+      writer.writerow(['seed', 'goal_error', 'solve_time_sec', 'solver_iters'])
+
+    visualizer = Visualizer()
+    visualizer.run()
+
+    method = None
+    for seed in range(_SEED.value, _SEED.value + _TESTS.value):
+      task = task_cls(seed)
+      if method is None:
+        method = method_cls(task)
+      elif hasattr(method, 'update'):
+        method.update(task)
+      else:
+        method = method_cls(task)
+
+      iters = _ITERATIONS.value
+
+      tic = time.perf_counter()
+
+      traj = method.solve(
+        traj_init={
+            'xs': task.xs_init,
+            'us': task.us_init,
+        },
+        iterations=iters,
+      )
+
+      solve_time = time.perf_counter() - tic
+
+      goal_error = support.state_cost(
+          task.Q, task.x_goal, traj['xs'][-1], task.quaternion_w_idxs
+      )
+
+      writer.writerow([seed, goal_error, solve_time, iters])
+      f.flush()
+
+      traj['info'] = f'''
+- Task: {task_cls.__name__}(seed={seed})
+  - goal error: {goal_error:.8f}
+- Method: {method_cls.__name__}
+  - solve time: {solve_time:.4f} sec
+  - iterations: {iters}
+'''
+      visualizer.update(task, traj)
+
+  if _TESTS.value == 1:
+    while True:
+      time.sleep(3600)
+
+
+if __name__ == '__main__':
+  app.run(_main)

mujodico-3.15.0/mujodico/examples/methods/__init__.py

@@ -0,0 +1,49 @@
+from mujodico.examples.methods.ipopt_contact_implicit_shooting import IpoptContactImplicitShooting
+from mujodico.examples.methods.ipopt_lifted_multiple_shooting import IpoptLiftedMultipleShooting
+from mujodico.examples.methods.ipopt_multiple_shooting import IpoptMultipleShooting
+from mujodico.trajopt._src.ipopt_multiple_shooting_solver import IpoptMultipleShootingSolver
+from mujodico.examples.methods.ipopt_single_shooting import IpoptSingleShooting
+from mujodico.examples.methods.ippas_lifted_multiple_shooting import IppasLiftedMultipleShooting
+from mujodico.examples.methods.transcription_lifted_multiple_shooting import LiftedMultipleShootingTranscription
+from mujodico.examples.methods.transcription_multiple_shooting import MultipleShootingTranscription
+from mujodico.examples.methods.transcription_single_shooting import SingleShootingTranscription
+
+
+__all__ = [
+    'IpoptContactImplicitShooting',
+    'IpoptLiftedMultipleShooting',
+    'IpoptMultipleShooting',
+    'IpoptMultipleShootingSolver',
+    'IpoptSingleShooting',
+    'IppasLiftedMultipleShooting',
+    'LiftedMultipleShootingTranscription',
+    'MultipleShootingTranscription',
+    'SingleShootingTranscription',
+    'get_method_names',
+    'parse_method_name',
+]
+
+
+_METHODS = [
+    IpoptContactImplicitShooting,
+    IpoptLiftedMultipleShooting,
+    IpoptMultipleShooting,
+    IpoptSingleShooting,
+    IppasLiftedMultipleShooting,
+]
+
+
+def get_method_names():
+  return [method.__name__ for method in _METHODS]
+
+
+def parse_method_name(name: str):
+  name = name.lower().strip().translate(str.maketrans('', '', ' -_'))
+
+  for method in _METHODS:
+    if name == method.__name__.lower():
+      return method
+
+  raise ValueError(
+      f'unknown method: value should be one of <{"|".join(get_method_names())}>'
+  )