| Name | mjinx JSON |
| Version |
1.0.0
JSON |
| download |
| home_page | None |
| Summary | Numerical Inverse Kinematics based on JAX + MJX |
| upload_time | 2024-10-26 14:01:48 |
| maintainer | None |
| docs_url | None |
| author | None |
| requires_python | >=3.10 |
| license | None |
| keywords |
|
| VCS |
|
| bugtrack_url |
|
| requirements |
No requirements were recorded.
|
| Travis-CI |
No Travis.
|
| coveralls test coverage |
No coveralls.
|
# MJINX
<!-- [](https://github.com/based-robotics/mjinx/actions)
[](https://github.com/based-robotics/mjinx/actions)
[](https://github.com/based-robotics/mjinx/actions)
[](https://pypi.org/project/jaxadi/)
[](https://pypistats.org/packages/jaxadi) -->
**Mjinx** is a library for auto-differentiable numerical inverse kinematics, powered by **JAX** and **Mujoco MJX**. The library was heavily inspired by the similar Pinocchio-based tool [pink](https://github.com/stephane-caron/pink/tree/main) and Mujoco-based analogue [mink](https://github.com/kevinzakka/mink/tree/main).
<!-- <div align="center">
<img src="img/local_ik_output.gif" style="width: 45%; max-width: 300px" />
<img src="img/go2_stance.gif" style="width: 45%; max-width: 300px" />
</div> -->
<p align="center">
<img src="img/local_ik_output.gif" style="width: 300px" />
<img src="img/go2_stance.gif" style="width: 300px" />
<img src="img/local_ik_input.gif" style="width: 300px"/>
</p>
## Key features
1. *Flexibility*. Each control problem is assembled via `Components`, which enforce desired behaviour or keeps system in a safety set.
2. *Different solution approaches*. `JAX` (i.e. it's efficient sampling and autodifferentiation) allows to implement variety of solvers, which might be more beneficial in different scenarios.
3. *Fully Jax-compatible*. Both optimal control problem and its solver are jax-compatible: jit-compilation and automatic vectorization are available for the whole problem.
4. *Convinience*. The functionality is nicely wrapped to make the interaction with it easier.
## Installation
The package is available in PyPI registry, and could be installed via `pip`:
```python
pip install mjinx
```
To run an examples or tests, please install the development version by running:
```python
pip install mjinx[dev]
```
## Usage
Here is the example of `mjinx` usage:
```python
from mujoco import mjx mjx
from mjinx.problem import Problem
# Initialize the robot model using MuJoCo
MJCF_PATH = "path_to_mjcf.xml"
mj_model = mj.MjModel.from_xml_path(MJCF_PATH)
mjx_model = mjx.put_model(mj_model)
# Create instance of the problem
problem = Problem(mjx_model)
# Add tasks to track desired behavior
frame_task = FrameTask("ee_task", cost=1, gain=20, body_name="link7")
problem.add_component(frame_task)
# Add barriers to keep robot in a safety set
joints_barrier = JointBarrier("jnt_range", gain=10)
problem.add_component(joitns_barrier)
# Initialize the solver
solver = LocalIKSolver(mjx_model)
# Initializing initial condition
q0 = np.zeros(7)
# Initialize solver data
solver_data = solver.init()
# jit-compiling solve and integrate
solve_jit = jax.jit(solver.solve)
integrate_jit = jax.jit(integrate, static_argnames=["dt"])
# === Control loop ===
for t in np.arange(0, 5, 1e-2):
# Changing problem and compiling it
frame_task.target_frame = np.array([0.1 * np.sin(t), 0.1 * np.cos(t), 0.1, 1, 0, 0,])
problem_data = problem.compile()
# Solving the instance of the problem
opt_solution, solver_data = solve_jit(q, solver_data, problem_data)
# Integrating
q = integrate_jit(
mjx_model,
q,
opt_solution.v_opt,
dt,
)
```
## Examples
The list of examples includes:
1. `Kuka iiwa` local inverse kinematics ([single item](examples/local_ik.py), [vmap over desired trajectory](examples/local_ik_vmapped_output.py))
2. `Kuka iiwa` global inverse kinematics ([single item](examples/global_ik.py), [vmap over desired trajectory](examples/global_ik_vmapped_output.py))
3. `Go2` [batched squats](examples/go2_squat.py) example
## Contributing
We are always open for the suggestions and contributions. For contribution guidelines, see the [CONTRIBUTING.md](CONTRIBUTING.md) file.
### TODO
The repostiory is under active development, the current plans before release are:
- [ ] Add examples for:
- [ ] Quadrotor
- [ ] Bipedal robot
- [ ] An MPPI example
- [ ] (?) Collaboration of several robots
- [ ] Add github pages
- [ ] Extend mathematical descriptions in docstrings
- [ ] Add potential fields example
## Acknowledgement
The repository was highly inspired by [`pink`](https://github.com/stephane-caron/pink) and [`mink`](https://github.com/kevinzakka/mink). Both authors, Stéphane Caron and Kevin Zakka, deeply inspire me to study robotics and contribute to the open source. Without them, this repository would not exist.
Some utility functions in this code are taken from source code of [`MuJoCo MJX`](https://github.com/google-deepmind/mujoco/tree/main/mjx). Apart from being a wonderful tool for batched computations and ML, the source code is compact yet readable and informative, and I encourage everyone to take a look at it to learn a little bit more about physical simulations, `jax`, and MuJoCo in general.
Raw data
{
"_id": null,
"home_page": null,
"name": "mjinx",
"maintainer": null,
"docs_url": null,
"requires_python": ">=3.10",
"maintainer_email": null,
"keywords": null,
"author": null,
"author_email": "Ivan Domrachev <domrachev10@gmail.com>, Simeon Nedelchev <simkaned@gmail.com>, Lev Kozlov <kozlov.l.a10@gmail.com>",
"download_url": null,
"platform": null,
"description": "# MJINX\n<!-- [](https://github.com/based-robotics/mjinx/actions)\n[](https://github.com/based-robotics/mjinx/actions)\n[](https://github.com/based-robotics/mjinx/actions)\n[](https://pypi.org/project/jaxadi/)\n[](https://pypistats.org/packages/jaxadi) -->\n\n\n**Mjinx** is a library for auto-differentiable numerical inverse kinematics, powered by **JAX** and **Mujoco MJX**. The library was heavily inspired by the similar Pinocchio-based tool [pink](https://github.com/stephane-caron/pink/tree/main) and Mujoco-based analogue [mink](https://github.com/kevinzakka/mink/tree/main).\n\n<!-- <div align=\"center\">\n <img src=\"img/local_ik_output.gif\" style=\"width: 45%; max-width: 300px\" />\n <img src=\"img/go2_stance.gif\" style=\"width: 45%; max-width: 300px\" /> \n</div> -->\n<p align=\"center\">\n <img src=\"img/local_ik_output.gif\" style=\"width: 300px\" />\n <img src=\"img/go2_stance.gif\" style=\"width: 300px\" /> \n <img src=\"img/local_ik_input.gif\" style=\"width: 300px\"/>\n</p>\n\n## Key features\n1. *Flexibility*. Each control problem is assembled via `Components`, which enforce desired behaviour or keeps system in a safety set. \n2. *Different solution approaches*. `JAX` (i.e. it's efficient sampling and autodifferentiation) allows to implement variety of solvers, which might be more beneficial in different scenarios.\n3. *Fully Jax-compatible*. Both optimal control problem and its solver are jax-compatible: jit-compilation and automatic vectorization are available for the whole problem.\n4. *Convinience*. The functionality is nicely wrapped to make the interaction with it easier.\n\n## Installation\nThe package is available in PyPI registry, and could be installed via `pip`:\n```python\npip install mjinx\n```\n\nTo run an examples or tests, please install the development version by running:\n```python\npip install mjinx[dev]\n```\n\n## Usage\nHere is the example of `mjinx` usage:\n\n```python\nfrom mujoco import mjx mjx\nfrom mjinx.problem import Problem\n\n# Initialize the robot model using MuJoCo\nMJCF_PATH = \"path_to_mjcf.xml\"\nmj_model = mj.MjModel.from_xml_path(MJCF_PATH)\nmjx_model = mjx.put_model(mj_model)\n\n# Create instance of the problem\nproblem = Problem(mjx_model)\n\n# Add tasks to track desired behavior\nframe_task = FrameTask(\"ee_task\", cost=1, gain=20, body_name=\"link7\")\nproblem.add_component(frame_task)\n\n# Add barriers to keep robot in a safety set\njoints_barrier = JointBarrier(\"jnt_range\", gain=10)\nproblem.add_component(joitns_barrier)\n\n# Initialize the solver\nsolver = LocalIKSolver(mjx_model)\n\n# Initializing initial condition\nq0 = np.zeros(7)\n\n# Initialize solver data\nsolver_data = solver.init()\n\n# jit-compiling solve and integrate \nsolve_jit = jax.jit(solver.solve)\nintegrate_jit = jax.jit(integrate, static_argnames=[\"dt\"])\n\n# === Control loop ===\nfor t in np.arange(0, 5, 1e-2):\n # Changing problem and compiling it\n frame_task.target_frame = np.array([0.1 * np.sin(t), 0.1 * np.cos(t), 0.1, 1, 0, 0,])\n problem_data = problem.compile()\n\n # Solving the instance of the problem\n opt_solution, solver_data = solve_jit(q, solver_data, problem_data)\n\n # Integrating\n q = integrate_jit(\n mjx_model,\n q,\n opt_solution.v_opt,\n dt,\n )\n```\n\n## Examples\nThe list of examples includes:\n 1. `Kuka iiwa` local inverse kinematics ([single item](examples/local_ik.py), [vmap over desired trajectory](examples/local_ik_vmapped_output.py))\n 2. `Kuka iiwa` global inverse kinematics ([single item](examples/global_ik.py), [vmap over desired trajectory](examples/global_ik_vmapped_output.py))\n 3. `Go2` [batched squats](examples/go2_squat.py) example\n \n\n## Contributing\nWe are always open for the suggestions and contributions. For contribution guidelines, see the [CONTRIBUTING.md](CONTRIBUTING.md) file. \n\n### TODO\nThe repostiory is under active development, the current plans before release are:\n- [ ] Add examples for:\n - [ ] Quadrotor\n - [ ] Bipedal robot\n - [ ] An MPPI example\n - [ ] (?) Collaboration of several robots\n- [ ] Add github pages\n - [ ] Extend mathematical descriptions in docstrings\n- [ ] Add potential fields example\n\n## Acknowledgement\nThe repository was highly inspired by [`pink`](https://github.com/stephane-caron/pink) and [`mink`](https://github.com/kevinzakka/mink). Both authors, St\u00e9phane Caron and Kevin Zakka, deeply inspire me to study robotics and contribute to the open source. Without them, this repository would not exist.\n\nSome utility functions in this code are taken from source code of [`MuJoCo MJX`](https://github.com/google-deepmind/mujoco/tree/main/mjx). Apart from being a wonderful tool for batched computations and ML, the source code is compact yet readable and informative, and I encourage everyone to take a look at it to learn a little bit more about physical simulations, `jax`, and MuJoCo in general. \n",
"bugtrack_url": null,
"license": null,
"summary": "Numerical Inverse Kinematics based on JAX + MJX",
"version": "1.0.0",
"project_urls": null,
"split_keywords": [],
"urls": [
{
"comment_text": "",
"digests": {
"blake2b_256": "1497068a2798d1430bbdd1e087be0790c3794826752499be0bc0c40e271e251d",
"md5": "681c9ceccbcfba865a74e9025cd77494",
"sha256": "1fa6f284225b9394088e79264699036d818a54fcd09ccc51c4c93db173938a20"
},
"downloads": -1,
"filename": "mjinx-1.0.0-py3-none-any.whl",
"has_sig": false,
"md5_digest": "681c9ceccbcfba865a74e9025cd77494",
"packagetype": "bdist_wheel",
"python_version": "py3",
"requires_python": ">=3.10",
"size": 50088,
"upload_time": "2024-10-26T14:01:48",
"upload_time_iso_8601": "2024-10-26T14:01:48.683620Z",
"url": "https://files.pythonhosted.org/packages/14/97/068a2798d1430bbdd1e087be0790c3794826752499be0bc0c40e271e251d/mjinx-1.0.0-py3-none-any.whl",
"yanked": false,
"yanked_reason": null
}
],
"upload_time": "2024-10-26 14:01:48",
"github": false,
"gitlab": false,
"bitbucket": false,
"codeberg": false,
"lcname": "mjinx"
}
