From 9decb7389f0a9d7577243cc799899048a7ea9bdd Mon Sep 17 00:00:00 2001 From: Baruch Tabanpour Date: Thu, 19 Oct 2023 10:37:27 -0700 Subject: [PATCH 01/74] Make humanoid default in colab. PiperOrigin-RevId: 574925299 Change-Id: I619438b4c61b882db27c3984fb8c061ec0c39e4d --- mjx/tutorial.ipynb | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mjx/tutorial.ipynb b/mjx/tutorial.ipynb index 8b7b78d8..234673e8 100644 --- a/mjx/tutorial.ipynb +++ b/mjx/tutorial.ipynb @@ -879,7 +879,7 @@ "source": [ "# instantiate the environment\n", "\n", - "env_name = 'barkour' # @param ['barkour', 'humanoid']\n", + "env_name = 'humanoid' # @param ['barkour', 'humanoid']\n", "env = envs.get_environment(env_name)\n", "\n", "# define the jit reset/step functions\n", From 7a802a1820c9192bb550c41176f31d994759e388 Mon Sep 17 00:00:00 2001 From: Baruch Tabanpour Date: Thu, 19 Oct 2023 10:43:55 -0700 Subject: [PATCH 02/74] PUBLIX: Fix typo in mjx doc. https://github.com/google-deepmind/mujoco/issues/1106 PiperOrigin-RevId: 574928021 Change-Id: I8f3f096ba71c41da1fdfdae8bd0f88f160d2678e --- doc/mjx.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/doc/mjx.rst b/doc/mjx.rst index 9189afe7..154409e5 100644 --- a/doc/mjx.rst +++ b/doc/mjx.rst @@ -94,7 +94,7 @@ Neither ``mjx.Model`` nor ``mjx.Data`` are meant to be constructed manually. An mjx_model = mjx.device_put(model) mjx_data = mjx.make_data(model) -Using ``mx.make_data`` may be preferable when constructing batched ``mjx.Data`` structures inside of a ``vmap``. +Using ``mjx.make_data`` may be preferable when constructing batched ``mjx.Data`` structures inside of a ``vmap``. .. _MjxFunctions: From 79e95895154bddd6d8c93752d18fbd8972d150d7 Mon Sep 17 00:00:00 2001 From: Baruch Tabanpour Date: Thu, 19 Oct 2023 11:06:56 -0700 Subject: [PATCH 03/74] Change test_box_box slightly. PiperOrigin-RevId: 574937448 Change-Id: If84c8e192b230fea0e858fdd12beceedaa280f30 --- mjx/mujoco/mjx/_src/collision_driver_test.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mjx/mujoco/mjx/_src/collision_driver_test.py b/mjx/mujoco/mjx/_src/collision_driver_test.py index 71b7b553..e31b95c8 100644 --- a/mjx/mujoco/mjx/_src/collision_driver_test.py +++ b/mjx/mujoco/mjx/_src/collision_driver_test.py @@ -28,7 +28,7 @@ from mujoco import mjx from mujoco.mjx._src.types import Contact from mujoco.mjx._src.types import Data from mujoco.mjx._src.types import Model -# pylint: emable=g-importing-member +# pylint: enable=g-importing-member import numpy as np @@ -289,7 +289,7 @@ class ConvexTest(absltest.TestCase): - + From 1178bad3395228f34341b15eb58d6de4b68a7e66 Mon Sep 17 00:00:00 2001 From: Saran Tunyasuvunakool Date: Thu, 19 Oct 2023 11:12:00 -0700 Subject: [PATCH 04/74] Run MJX tests on GitHub Actions. PiperOrigin-RevId: 574939328 Change-Id: Ic6e5650ef6a7df8c94a189206fddf06bd9c4eb97 --- .github/workflows/build.yml | 23 ++++++++++++++++++++++- 1 file changed, 22 insertions(+), 1 deletion(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 84646f60..7abd4c7e 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -219,7 +219,7 @@ jobs: working-directory: python/dist run: > source ${{ matrix.tmpdir }}/venv/bin/activate && - pip install --no-deps --no-index mujoco-*.whl + pip install --no-index mujoco-*.whl - name: Test Python bindings if: ${{ runner.os != 'Windows' }} shell: bash @@ -228,6 +228,27 @@ jobs: run: > source ${{ matrix.tmpdir }}/venv/bin/activate && pytest -v --pyargs mujoco + - name: Package MJX + if: ${{ runner.os != 'Windows' }} + shell: bash + working-directory: mjx + run: + source ${{ matrix.tmpdir }}/venv/bin/activate && + python -m build . + - name: Install MJX + if: ${{ runner.os != 'Windows' }} + shell: bash + working-directory: mjx + run: + source ${{ matrix.tmpdir }}/venv/bin/activate && + pip install --require-hashes -r requirements.txt && + pip install --no-index dist/mujoco_mjx-*.whl + - name: Test MJX + if: ${{ runner.os != 'Windows' }} + shell: bash + run: + source ${{ matrix.tmpdir }}/venv/bin/activate && + pytest -n auto -v --pyargs mujoco.mjx - name: Notify team chat shell: bash env: From dc8bac2f4dd7353208fa92d40bb3fafb7a82eb0a Mon Sep 17 00:00:00 2001 From: Baruch Tabanpour Date: Thu, 19 Oct 2023 11:42:11 -0700 Subject: [PATCH 05/74] Fix supported equality constraint in doc. https://github.com/google-deepmind/mujoco/issues/1105 PiperOrigin-RevId: 574950690 Change-Id: I1c029e9b31235fdcaa8e8f35b5117fec373fa25f --- doc/changelog.rst | 6 ++++++ doc/mjx.rst | 4 ++++ 2 files changed, 10 insertions(+) diff --git a/doc/changelog.rst b/doc/changelog.rst index 151c6ebf..c9f62f1a 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -10,6 +10,12 @@ Bug fixes 1. Fix in simulate: correct handling of "Pause update", "Fullscreen" and "VSync" buttons. +MJX Bug fixes +^^^^^^^^^ + +1. Fix typos and supported fields in the docs (fixes :github:issue:`1105` and :github:issue:`1106`). + + Version 3.0.0 (October 18, 2023) -------------------------------- diff --git a/doc/mjx.rst b/doc/mjx.rst index 154409e5..92566ddb 100644 --- a/doc/mjx.rst +++ b/doc/mjx.rst @@ -187,6 +187,8 @@ The following features are **fully supported** in MJX: - ``PLANE``, ``SPHERE``, ``CAPSULE``, ``BOX``, ``MESH`` * - :ref:`Constraint ` - ``EQUALITY``, ``FRICTION_DOF``, ``LIMIT_JOINT``, ``CONTACT_PYRAMIDAL`` + * - :ref:`Equality ` + - ``CONNECT``, ``WELD`` * - :ref:`Integrator ` - ``EULER``, ``RK4`` * - :ref:`Cone ` @@ -226,6 +228,8 @@ The following features are **in development** and coming soon: - :ref:`flEllipsoid` * - :ref:`Tendons ` - :ref:`Spatial `, :ref:`Fixed ` + * - :ref:`Equality ` + - ``JOINT``, ``TENDON`` The following features are **unsupported**: From dd841c92831fd470c5570561efe07036a260efd1 Mon Sep 17 00:00:00 2001 From: Baruch Tabanpour Date: Thu, 19 Oct 2023 15:39:15 -0700 Subject: [PATCH 06/74] Fix bug in joint range for MJX. PiperOrigin-RevId: 575023489 Change-Id: I717333064ed4d6971a841bec60e8afd3019120a7 --- doc/changelog.rst | 2 +- mjx/mujoco/mjx/_src/constraint.py | 8 ++++- mjx/mujoco/mjx/_src/constraint_test.py | 34 ++++++++++++++++++++ mjx/mujoco/mjx/test_data/cherry_pendulum.xml | 2 +- 4 files changed, 43 insertions(+), 3 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index c9f62f1a..e212eb61 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -14,7 +14,7 @@ MJX Bug fixes ^^^^^^^^^ 1. Fix typos and supported fields in the docs (fixes :github:issue:`1105` and :github:issue:`1106`). - +2. Fix bug where mixed `jnt_limited` joints are not being constrained correctly. Version 3.0.0 (October 18, 2023) -------------------------------- diff --git a/mjx/mujoco/mjx/_src/constraint.py b/mjx/mujoco/mjx/_src/constraint.py index bd0bbaa6..3f020198 100644 --- a/mjx/mujoco/mjx/_src/constraint.py +++ b/mjx/mujoco/mjx/_src/constraint.py @@ -228,13 +228,19 @@ def _instantiate_limit(m: Model, d: Data) -> _Efc: return jp.stack(js), jp.stack(rs), jp.stack(arefs) + jnt_range = jp.where( + m.jnt_limited[:, None], + m.jnt_range, + jp.array([-jp.inf, jp.inf]), + ) + j, r, aref = scan.flat( m, fn, 'jjjjjqvv', 'jjj', m.jnt_type, - m.jnt_range, + jnt_range, m.jnt_solref, m.jnt_solimp, m.jnt_margin, diff --git a/mjx/mujoco/mjx/_src/constraint_test.py b/mjx/mujoco/mjx/_src/constraint_test.py index ce289df7..8bc77042 100644 --- a/mjx/mujoco/mjx/_src/constraint_test.py +++ b/mjx/mujoco/mjx/_src/constraint_test.py @@ -24,6 +24,7 @@ from mujoco.mjx._src import constraint from mujoco.mjx._src import test_util # pylint: disable=g-importing-member from mujoco.mjx._src.types import DisableBit +from mujoco.mjx._src.types import SolverType # pylint: enable=g-importing-member import numpy as np @@ -79,6 +80,39 @@ class ConstraintTest(parameterized.TestCase): fname, ) + _JNT_RANGE = """ + + + + + + + + + + + + + """ + + def test_jnt_range(self): + """Tests that mixed joint ranges are respected.""" + m = mujoco.MjModel.from_xml_string(self._JNT_RANGE) + m.opt.solver = SolverType.CG.value + d = mujoco.MjData(m) + d.qpos = np.array([2.0, 15.0]) + + mx = mjx.device_put(m) + dx = mjx.device_put(d) + efc = jax.jit(constraint._instantiate_limit)(mx, dx) + + # first joint is outside the joint range + np.testing.assert_array_almost_equal(efc.J[0, 0], -1.0) + + # second joint does not hit joint range + np.testing.assert_array_almost_equal(efc.aref[1], 0.0) + def test_disable_refsafe(self): m = test_util.load_test_file('ant.xml') diff --git a/mjx/mujoco/mjx/test_data/cherry_pendulum.xml b/mjx/mujoco/mjx/test_data/cherry_pendulum.xml index bfe2db9b..ea9082de 100644 --- a/mjx/mujoco/mjx/test_data/cherry_pendulum.xml +++ b/mjx/mujoco/mjx/test_data/cherry_pendulum.xml @@ -5,7 +5,7 @@ - + From e4a9f535a4fe05c2120ae9cd7afb17665f917e6b Mon Sep 17 00:00:00 2001 From: Baruch Tabanpour Date: Thu, 19 Oct 2023 17:28:50 -0700 Subject: [PATCH 07/74] Make MJX device_put type validation more verbose. PiperOrigin-RevId: 575050350 Change-Id: I7cab4ee54d462b61e9478dbd58d99bb9fb662daf --- doc/changelog.rst | 1 + mjx/mujoco/mjx/_src/device.py | 5 +++-- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index e212eb61..e64a7300 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -15,6 +15,7 @@ MJX Bug fixes 1. Fix typos and supported fields in the docs (fixes :github:issue:`1105` and :github:issue:`1106`). 2. Fix bug where mixed `jnt_limited` joints are not being constrained correctly. +3. Make `device_put` type validation more verbose (fixed :github:issue:`1113`). Version 3.0.0 (October 18, 2023) -------------------------------- diff --git a/mjx/mujoco/mjx/_src/device.py b/mjx/mujoco/mjx/_src/device.py index 3c3c213f..a917e7f9 100644 --- a/mjx/mujoco/mjx/_src/device.py +++ b/mjx/mujoco/mjx/_src/device.py @@ -116,8 +116,9 @@ def _validate(m: mujoco.MjModel): typs = set(val) if isinstance(val, Iterable) else {val} unsupported_typs = typs - set(_TYPE_MAP[mj_type]) - if unsupported_typs: - raise NotImplementedError(f'{unsupported_typs} not implemented.') + unsupported = [mj_type(t) for t in unsupported_typs] + if unsupported: + raise NotImplementedError(f'{unsupported} not implemented.') # check condim if any(dim != 3 for dim in m.geom_condim) or any( From 01a932e052f599ca0d26d8060d759ea3965c5bea Mon Sep 17 00:00:00 2001 From: Alessio Quaglino Date: Fri, 20 Oct 2023 02:48:46 -0700 Subject: [PATCH 08/74] Fix flex constraint Jacobian for edges belonging to a single body. PiperOrigin-RevId: 575161094 Change-Id: If4153bcb751cdd6dbdddeb08672f7107c7c5bbc2 --- doc/includes/references.h | 3 +- include/mujoco/mjmodel.h | 3 +- include/mujoco/mjxmacro.h | 1 + introspect/structs.py | 9 ++++- src/engine/engine_core_constraint.c | 52 ++++++++++++++++++---------- src/engine/engine_core_smooth.c | 12 +++++-- src/engine/engine_derivative.c | 11 ++++-- src/engine/engine_passive.c | 7 +++- src/engine/engine_sensor.c | 12 ++++--- src/engine/engine_setconst.c | 7 +++- src/user/user_model.cc | 17 +++++---- unity/Runtime/Bindings/MjBindings.cs | 1 + 12 files changed, 97 insertions(+), 38 deletions(-) diff --git a/doc/includes/references.h b/doc/includes/references.h index 3f98f313..1ffd8c4b 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -944,7 +944,7 @@ struct mjModel_ { int* body_treeid; // id of body's kinematic tree; -1: static (nbody x 1) int* body_geomnum; // number of geoms (nbody x 1) int* body_geomadr; // start addr of geoms; -1: no geoms (nbody x 1) - mjtByte* body_simple; // 1: diagonal M; 2: diag M, no rotations (nbody x 1) + mjtByte* body_simple; // 1: diag M; 2: diag M, sliders only (nbody x 1) mjtByte* body_sameframe; // inertial frame is same as body frame (nbody x 1) mjtNum* body_pos; // position offset rel. to parent body (nbody x 3) mjtNum* body_quat; // orientation offset rel. to parent body (nbody x 4) @@ -1122,6 +1122,7 @@ struct mjModel_ { mjtNum* flex_edgedamping; // edge damping (nflex x 1) mjtByte* flex_edgeequality; // is edge equality constraint defined (nflex x 1) mjtByte* flex_rigid; // are all verices in the same body (nflex x 1) + mjtByte* flexedge_rigid; // are both edge vertices in same body (nflexedge x 1) mjtByte* flex_centered; // are all vertex coordinates (0,0,0) (nflex x 1) mjtByte* flex_flatskin; // render flex skin with flat shading (nflex x 1) int* flex_bvhadr; // address of bvh root; -1: no bvh (nflex x 1) diff --git a/include/mujoco/mjmodel.h b/include/mujoco/mjmodel.h index c84c9e3e..a781f7fa 100644 --- a/include/mujoco/mjmodel.h +++ b/include/mujoco/mjmodel.h @@ -661,7 +661,7 @@ struct mjModel_ { int* body_treeid; // id of body's kinematic tree; -1: static (nbody x 1) int* body_geomnum; // number of geoms (nbody x 1) int* body_geomadr; // start addr of geoms; -1: no geoms (nbody x 1) - mjtByte* body_simple; // 1: diagonal M; 2: diag M, no rotations (nbody x 1) + mjtByte* body_simple; // 1: diag M; 2: diag M, sliders only (nbody x 1) mjtByte* body_sameframe; // inertial frame is same as body frame (nbody x 1) mjtNum* body_pos; // position offset rel. to parent body (nbody x 3) mjtNum* body_quat; // orientation offset rel. to parent body (nbody x 4) @@ -839,6 +839,7 @@ struct mjModel_ { mjtNum* flex_edgedamping; // edge damping (nflex x 1) mjtByte* flex_edgeequality; // is edge equality constraint defined (nflex x 1) mjtByte* flex_rigid; // are all verices in the same body (nflex x 1) + mjtByte* flexedge_rigid; // are both edge vertices in same body (nflexedge x 1) mjtByte* flex_centered; // are all vertex coordinates (0,0,0) (nflex x 1) mjtByte* flex_flatskin; // render flex skin with flat shading (nflex x 1) int* flex_bvhadr; // address of bvh root; -1: no bvh (nflex x 1) diff --git a/include/mujoco/mjxmacro.h b/include/mujoco/mjxmacro.h index 01b310ac..eca00e80 100644 --- a/include/mujoco/mjxmacro.h +++ b/include/mujoco/mjxmacro.h @@ -345,6 +345,7 @@ X ( mjtNum, flex_edgedamping, nflex, 1 ) \ X ( mjtByte, flex_edgeequality, nflex, 1 ) \ X ( mjtByte, flex_rigid, nflex, 1 ) \ + X ( mjtByte, flexedge_rigid, nflexedge, 1 ) \ X ( mjtByte, flex_centered, nflex, 1 ) \ XMJV( mjtByte, flex_flatskin, nflex, 1 ) \ XMJV( int, flex_bvhadr, nflex, 1 ) \ diff --git a/introspect/structs.py b/introspect/structs.py index 639600cb..17579509 100644 --- a/introspect/structs.py +++ b/introspect/structs.py @@ -1311,7 +1311,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([ type=PointerType( inner_type=ValueType(name='mjtByte'), ), - doc='1: diagonal M; 2: diag M, no rotations (nbody x 1)', + doc='1: diag M; 2: diag M, sliders only (nbody x 1)', ), StructFieldDecl( name='body_sameframe', @@ -2426,6 +2426,13 @@ STRUCTS: Mapping[str, StructDecl] = dict([ ), doc='are all verices in the same body (nflex x 1)', ), + StructFieldDecl( + name='flexedge_rigid', + type=PointerType( + inner_type=ValueType(name='mjtByte'), + ), + doc='are both edge vertices in same body (nflexedge x 1)', + ), StructFieldDecl( name='flex_centered', type=PointerType( diff --git a/src/engine/engine_core_constraint.c b/src/engine/engine_core_constraint.c index d6f723f6..dbe268d9 100644 --- a/src/engine/engine_core_constraint.c +++ b/src/engine/engine_core_constraint.c @@ -663,9 +663,13 @@ void mj_instantiateEquality(const mjModel* m, mjData* d) { case mjEQ_FLEX: flex_edgeadr = m->flex_edgeadr[id[0]]; flex_edgenum = m->flex_edgenum[id[0]]; - - // add one constraint per edge + // add one constraint per non-rigid edge for (int e=flex_edgeadr; e < flex_edgeadr+flex_edgenum; e++) { + // skip rigid + if (m->flexedge_rigid[e]) { + continue; + } + // position error cpos[0] = d->flexedge_length[e] - m->flexedge_length0[e]; @@ -1057,7 +1061,7 @@ void mj_instantiateContact(const mjModel* m, mjData* d) { // compute diagApprox void mj_diagApprox(const mjModel* m, mjData* d) { - int id, dim, b1, b2, weldcnt = 0, edgecnt = 0; + int id, dim, b1, b2, f, weldcnt = 0; int nefc = d->nefc; mjtNum tran, rot, fri, *dA = d->efc_diagApprox; mjContact* con = NULL; @@ -1067,11 +1071,6 @@ void mj_diagApprox(const mjModel* m, mjData* d) { // get constraint id id = d->efc_id[i]; - // clear edge counter - if (d->efc_type[i] != mjEQ_FLEX) { - edgecnt = 0; - } - // process according to constraint type switch ((mjtConstraint) d->efc_type[i]) { case mjCNSTR_EQUALITY: @@ -1108,8 +1107,18 @@ void mj_diagApprox(const mjModel* m, mjData* d) { break; case mjEQ_FLEX: - dA[i] = m->flexedge_invweight0[m->flex_edgeadr[m->eq_obj1id[id]] + edgecnt]; - edgecnt++; + // process all non-rigid edges for this flex + f = m->eq_obj1id[id]; + int flex_edgeadr = m->flex_edgeadr[f]; + int flex_edgenum = m->flex_edgenum[f]; + for (int e=flex_edgeadr; eflexedge_rigid[e]) { + dA[i++] = m->flexedge_invweight0[e]; + } + } + + // adjust constraint counter + i--; break; default: @@ -1642,21 +1651,26 @@ static int mj_ne(const mjModel* m, mjData* d, int* nnz) { break; case mjEQ_FLEX: - size = m->flex_edgenum[id[0]]; - if (!nnz) { - break; - } - flex_edgeadr = m->flex_edgeadr[id[0]]; flex_edgenum = m->flex_edgenum[id[0]]; + // init with all edges, subract rigid later + size = flex_edgenum; + // process edges of this flex for (int e=flex_edgeadr; e < flex_edgeadr+flex_edgenum; e++) { - int b1 = m->flex_vertbodyid[m->flex_vertadr[id[0]] + m->flex_edge[2*e]]; - int b2 = m->flex_vertbodyid[m->flex_vertadr[id[0]] + m->flex_edge[2*e+1]]; + // rigid: reduce size and skip + if (m->flexedge_rigid[e]) { + size--; + continue; + } - // accumulate NV - NV += mj_jacDifPairCount(m, chain, b1, b2, issparse); + // accumulate NV if needed + if (nnz) { + int b1 = m->flex_vertbodyid[m->flex_vertadr[id[0]] + m->flex_edge[2*e]]; + int b2 = m->flex_vertbodyid[m->flex_vertadr[id[0]] + m->flex_edge[2*e+1]]; + NV += mj_jacDifPairCount(m, chain, b1, b2, issparse); + } } break; diff --git a/src/engine/engine_core_smooth.c b/src/engine/engine_core_smooth.c index 260c41af..af8b10e6 100644 --- a/src/engine/engine_core_smooth.c +++ b/src/engine/engine_core_smooth.c @@ -1857,8 +1857,16 @@ void mj_rnePostConstraint(const mjModel* m, mjData* d) { break; case mjEQ_FLEX: - // increment edgenum rows - i += m->flex_edgenum[m->eq_obj1id[id]]; + // increment with number of non-rigid edges + k = m->eq_obj1id[id]; + int flex_edgeadr = m->flex_edgeadr[k]; + int flex_edgenum = m->flex_edgenum[k]; + + for (int e=flex_edgeadr; e < flex_edgeadr+flex_edgenum; e++) { + if (!m->flexedge_rigid[e]) { + i++; + } + } break; default: diff --git a/src/engine/engine_derivative.c b/src/engine/engine_derivative.c index 8414aec0..c0c0d5fe 100644 --- a/src/engine/engine_derivative.c +++ b/src/engine/engine_derivative.c @@ -1474,9 +1474,16 @@ void mjd_passive_vel(const mjModel* m, mjData* d) { for (int f=0; f < m->nflex; f++) { if (!m->flex_rigid[f] && m->flex_edgedamping[f]) { mjtNum B = -m->flex_edgedamping[f]; + int flex_edgeadr = m->flex_edgeadr[f]; + int flex_edgenum = m->flex_edgenum[f]; + + // process non-rigid edges of this flex + for (int e=flex_edgeadr; e < flex_edgeadr+flex_edgenum; e++) { + // skip rigid + if (m->flexedge_rigid[e]) { + continue; + } - // process edges of this flex - for (int e=m->flex_edgeadr[f]; e < m->flex_edgeadr[f]+m->flex_edgenum[f]; e++) { // add sparse or dense if (mj_isSparse(m)) { addJTBJSparse(m, d, d->flexedge_J, &B, 1, e, diff --git a/src/engine/engine_passive.c b/src/engine/engine_passive.c index 3da589fe..17a09d32 100644 --- a/src/engine/engine_passive.c +++ b/src/engine/engine_passive.c @@ -106,9 +106,14 @@ void mj_passive(const mjModel* m, mjData* d) { continue; } - // process edges of this flex (global edge index) + // process non-rigid edges of this flex (global edge index) int edgeend = m->flex_edgeadr[f] + m->flex_edgenum[f]; for (int e=m->flex_edgeadr[f]; e < edgeend; e++) { + // skip rigid + if (m->flexedge_rigid[e]) { + continue; + } + // compute spring-damper force along edge frc = stiffness * (m->flexedge_length0[e] - d->flexedge_length[e]) - damping * d->flexedge_velocity[e]; diff --git a/src/engine/engine_sensor.c b/src/engine/engine_sensor.c index 29d8b654..28ecb022 100644 --- a/src/engine/engine_sensor.c +++ b/src/engine/engine_sensor.c @@ -956,10 +956,14 @@ void mj_energyPos(const mjModel* m, mjData* d) { continue; } - // process edges of this flex - for (int e=m->flex_edgeadr[i]; e < m->flex_edgeadr[i]+m->flex_edgenum[i]; e++) { - mjtNum displacement = m->flexedge_length0[e] - d->flexedge_length[e]; - d->energy[0] += 0.5*stiffness*displacement*displacement; + // process non-rigid edges of this flex + int flex_edgeadr = m->flex_edgeadr[i]; + int flex_edgenum = m->flex_edgenum[i]; + for (int e=flex_edgeadr; e < flex_edgeadr+flex_edgenum; e++) { + if (!m->flexedge_rigid[e]) { + mjtNum displacement = m->flexedge_length0[e] - d->flexedge_length[e]; + d->energy[0] += 0.5*stiffness*displacement*displacement; + }; } } } diff --git a/src/engine/engine_setconst.c b/src/engine/engine_setconst.c index 0e9b5b1f..ff136d18 100644 --- a/src/engine/engine_setconst.c +++ b/src/engine/engine_setconst.c @@ -206,8 +206,13 @@ static void set0(mjModel* m, mjData* d) { int b1 = m->flex_vertbodyid[m->flex_vertadr[f] + m->flex_edge[2*i]]; int b2 = m->flex_vertbodyid[m->flex_vertadr[f] + m->flex_edge[2*i+1]]; + // rigid edge: set to 0 + if (m->flexedge_rigid[i]) { + m->flexedge_invweight0[i] = 0; + } + // accelerate edges that connect simple bodies with no rotations - if (m->body_simple[b1] == 2 && m->body_simple[b2] == 2) { + else if (m->body_simple[b1] == 2 && m->body_simple[b2] == 2) { m->flexedge_invweight0[i] = (1/m->body_mass[b1] + 1/m->body_mass[b2])/2; } diff --git a/src/user/user_model.cc b/src/user/user_model.cc index 657fc30a..855aa400 100644 --- a/src/user/user_model.cc +++ b/src/user/user_model.cc @@ -1569,8 +1569,8 @@ void mjCModel::CopyTree(mjModel* m) { qposadr += nPOS[pj->type]; } - // simple body with no rotational dofs: promote to simple level 2 - if (m->body_simple[i]) { + // simple body with sliders and no rotational dofs: promote to simple level 2 + if (m->body_simple[i] && m->body_dofnum[i]) { m->body_simple[i] = 2; for (int j=0; j<(int)pb->joints.size(); j++) { if (pb->joints[j]->type!=mjJNT_SLIDE) { @@ -1979,10 +1979,15 @@ void mjCModel::CopyObjects(mjModel* m) { memcpy(m->flex_vertbodyid + vert_adr, pfl->vertbodyid.data(), pfl->nvert*sizeof(int)); } - // convert edge pairs to int array - for (int i=0; inedge; i++) { - m->flex_edge[2*(edge_adr+i)] = pfl->edge[i].first; - m->flex_edge[2*(edge_adr+i)+1] = pfl->edge[i].second; + // convert edge pairs to int array, set edge rigid + for (int k=0; knedge; k++) { + m->flex_edge[2*(edge_adr+k)] = pfl->edge[k].first; + m->flex_edge[2*(edge_adr+k)+1] = pfl->edge[k].second; + + // check if vertex body weldids are the same + int b1 = pfl->vertbodyid[pfl->edge[k].first]; + int b2 = pfl->vertbodyid[pfl->edge[k].second]; + m->flexedge_rigid[edge_adr+k] = (bodies[b1]->weldid == bodies[b2]->weldid); } // advance counters diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs index f12f6fbb..a761f964 100644 --- a/unity/Runtime/Bindings/MjBindings.cs +++ b/unity/Runtime/Bindings/MjBindings.cs @@ -5337,6 +5337,7 @@ public unsafe struct mjModel_ { public double* flex_edgedamping; public byte* flex_edgeequality; public byte* flex_rigid; + public byte* flexedge_rigid; public byte* flex_centered; public byte* flex_flatskin; public int* flex_bvhadr; From ebe081f27e8e96083fee17f5e0d0e79cbedd07f6 Mon Sep 17 00:00:00 2001 From: Baruch Tabanpour Date: Fri, 20 Oct 2023 09:50:31 -0700 Subject: [PATCH 09/74] Add contact disable flag to collision driver. PiperOrigin-RevId: 575245746 Change-Id: If2ea8a5de6f6c751729c9b8dd5efc42e7bdd51ea --- mjx/mujoco/mjx/_src/collision_driver.py | 27 +++++++++++++++----- mjx/mujoco/mjx/_src/collision_driver_test.py | 26 +++++++++++++++++++ 2 files changed, 47 insertions(+), 6 deletions(-) diff --git a/mjx/mujoco/mjx/_src/collision_driver.py b/mjx/mujoco/mjx/_src/collision_driver.py index b5f58130..dde3fe92 100644 --- a/mjx/mujoco/mjx/_src/collision_driver.py +++ b/mjx/mujoco/mjx/_src/collision_driver.py @@ -336,16 +336,28 @@ def collision_candidates(m: Union[Model, mujoco.MjModel]) -> CandidateSet: def ncon(m: Model) -> int: """Returns the number of contacts computed in MJX given a model.""" + if m.opt.disableflags & DisableBit.CONTACT: + return 0 + candidates = collision_candidates(m) max_count = _max_contact_points(m) - count = sum([ - len(v) * get_collision_fn(k[0:2]).ncon for k, v in candidates.items() # pytype: disable=attribute-error - ]) + + count = 0 + for k, v in candidates.items(): + fn = get_collision_fn(k[0:2]) + if fn is None: + continue + count += len(v) * fn.ncon # pytype: disable=attribute-error + return min(max_count, count) if max_count > -1 else count def collision(m: Model, d: Data) -> Data: """Collides geometries.""" + ncon_ = ncon(m) + if ncon_ == 0: + return d.replace(contact=Contact.zero(), ncon=0) + candidate_set = collision_candidates(m) contacts = [] @@ -354,7 +366,7 @@ def collision(m: Model, d: Data) -> Data: contacts.append(_collide_geoms(m, d, geom_types, candidates)) if not contacts: - return d.replace(contact=Contact.zero(), ncon=0) + raise RuntimeError('No contacts found.') contact = jax.tree_map(lambda *x: jp.concatenate(x), *contacts) @@ -364,10 +376,13 @@ def collision(m: Model, d: Data) -> Data: _, idx = jax.lax.top_k(-contact.dist, k=max_contact_points) contact = jax.tree_map(lambda x, idx=idx: jp.take(x, idx, axis=0), contact) - ncon_ = contact.dist.shape[0] + if ncon_ != contact.dist.shape[0]: + raise RuntimeError('Number of contacts does not match ncon.') + + # TODO(robotics-simulation): move this logic to device_put ns = d.ne + d.nf + d.nl + contact = contact.replace(efc_address=np.arange(ns, ns + ncon_ * 4, 4)) # TODO(robotics-simulation): add support for other friction dimensions - contact = contact.replace(efc_address=np.arange(ns, ns + d.ncon * 4, 4)) contact = contact.replace(dim=3 * np.ones(ncon_, dtype=np.int32)) return d.replace(contact=contact, ncon=ncon_) diff --git a/mjx/mujoco/mjx/_src/collision_driver_test.py b/mjx/mujoco/mjx/_src/collision_driver_test.py index e31b95c8..b735e014 100644 --- a/mjx/mujoco/mjx/_src/collision_driver_test.py +++ b/mjx/mujoco/mjx/_src/collision_driver_test.py @@ -24,9 +24,12 @@ import jax import jax.numpy as jp import mujoco from mujoco import mjx +from mujoco.mjx._src import collision_driver +from mujoco.mjx._src import test_util # pylint: disable=g-importing-member from mujoco.mjx._src.types import Contact from mujoco.mjx._src.types import Data +from mujoco.mjx._src.types import DisableBit from mujoco.mjx._src.types import Model # pylint: enable=g-importing-member import numpy as np @@ -450,6 +453,29 @@ class BodyPairFilterTest(absltest.TestCase): self.assertEqual(dx.contact.pos.shape[0], 1) +class NconTest(parameterized.TestCase): + """Tests ncon.""" + + def test_ncon(self): + m = test_util.load_test_file('ant.xml') + d = mujoco.MjData(m) + d.qpos[2] = 0.0 + + mx = mjx.device_put(m) + ncon = collision_driver.ncon(mx) + self.assertEqual(ncon, 4) + + def test_disable_contact(self): + m = test_util.load_test_file('ant.xml') + d = mujoco.MjData(m) + d.qpos[2] = 0.0 + + m.opt.disableflags = m.opt.disableflags | DisableBit.CONTACT + mx = mjx.device_put(m) + ncon = collision_driver.ncon(mx) + self.assertEqual(ncon, 0) + + class TopKContactTest(absltest.TestCase): """Tests top-k contacts.""" From e57b34303a13195bf5fe87d37cb778321130a229 Mon Sep 17 00:00:00 2001 From: Baruch Tabanpour Date: Fri, 20 Oct 2023 09:53:09 -0700 Subject: [PATCH 10/74] Reformat notebook and add back sim2sim. PiperOrigin-RevId: 575246433 Change-Id: I8627916e7dedd5cf64dab6e0151f41c3153e7a68 --- mjx/tutorial.ipynb | 679 +++++++++++++++++++++++++++++---------------- 1 file changed, 433 insertions(+), 246 deletions(-) diff --git a/mjx/tutorial.ipynb b/mjx/tutorial.ipynb index 234673e8..83e1bd30 100644 --- a/mjx/tutorial.ipynb +++ b/mjx/tutorial.ipynb @@ -187,11 +187,9 @@ }, "source": [ "# Training a Policy with MJX\n", - "MJX is an implementation of MuJoCo written in [jax](https://jax.readthedocs.io/en/latest/index.html), enabling large batch training on GPU/TPU. In this notebook, we demonstrate training RL policies with MJX.\n", + "MJX is an implementation of MuJoCo written in [JAX](https://jax.readthedocs.io/en/latest/index.html), enabling large batch training on GPU/TPU. In this notebook, we demonstrate how to train RL policies with MJX.\n", "\n", - "First, we implement an environment `State` so that we can plug into the [brax](https://github.com/google/brax) environment API. `State` holds the observation, reward, metrics, and etc. Notably `State.pipeline_state` holds a `mjx.Data` object, which is analogous to `mjData` in MuJoCo.\n", - "\n", - "`MjxEnv` is an implementation of the `brax.envs.base.Env` class that initializes `mjx.Model` and `mjx.Data` objects. Inheriting from `brax.envs.base.Env` allows us to use the training agents implemented in brax. Notice that `MjxEnv` calls `mjx.step` for every `pipeline_step`, which is analgous to `mujoco.mj_step`." + "First, we implement an environment `State` so that we can plug into the [Brax](https://github.com/google/brax) environment API. `State` holds the observation, reward, metrics, and environment info. Notably `State.pipeline_state` holds a `mjx.Data` object, which is analogous to `mjData` in MuJoCo.\n" ] }, { @@ -202,7 +200,7 @@ }, "outputs": [], "source": [ - "#@title State and MjxEnv\n", + "#@title State\n", "\n", "@struct.dataclass\n", "class State(Base):\n", @@ -223,8 +221,30 @@ " reward: jax.Array\n", " done: jax.Array\n", " metrics: Dict[str, jax.Array] = struct.field(default_factory=dict)\n", - " info: Dict[str, Any] = struct.field(default_factory=dict)\n", + " info: Dict[str, Any] = struct.field(default_factory=dict)\n" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "acpXtDLNXLV9" + }, + "source": [ "\n", + "Next, we implement `MjxEnv`, an environment class we'll use through the notebook. `MjxEnv` initializes a `mjx.Model` and `mjx.Data` object. Notice that `MjxEnv` calls `mjx.step` for every `pipeline_step`, which is analgous to `mujoco.mj_step`.\n", + "\n", + "`MjxEnv` also inherits from `brax.envs.base.Env` which allows us to use the training agents implemented in brax." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "ccujYeJ5XOhx" + }, + "outputs": [], + "source": [ + "#@title MjxEnv\n", "\n", "class MjxEnv(Env):\n", " \"\"\"API for driving an MJX system for training and inference in brax.\"\"\"\n", @@ -305,7 +325,7 @@ "id": "iPlFu4CiIgBN" }, "source": [ - "Now we can define environment implementations for Humanoid and the [Barkour v0 Quadruped](https://github.com/google-deepmind/mujoco_menagerie/tree/main/google_barkour_v0). The environments define the reward and stepping logic. `reset` initializes a `State`, and `step` steps through the physics step and reward logic." + "Finally we can implement a real environment. We choose to first implement the Humanoid environment. Notice that `reset` initializes a `State`, and `step` steps through the physics step and reward logic. The reward and stepping logic train the Humanoid to run forwards." ] }, { @@ -444,6 +464,372 @@ "envs.register_environment('humanoid', Humanoid)" ] }, + { + "cell_type": "markdown", + "metadata": { + "id": "P1K6IznI2y83" + }, + "source": [ + "## Visualize a Rollout\n", + "\n", + "Let's instantiate the environment and visualize a short rollout.\n", + "\n", + "NOTE: Since episodes terminates early if the torso is below the healthy z-range, the only relevant contacts for this task are between the feet and the plane. We turn off other contacts." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "EhKLFK54C1CH" + }, + "outputs": [], + "source": [ + "# instantiate the environment\n", + "env_name = 'humanoid'\n", + "env = envs.get_environment(env_name)\n", + "\n", + "# define the jit reset/step functions\n", + "jit_reset = jax.jit(env.reset)\n", + "jit_step = jax.jit(env.step)\n", + "\n", + "# instantiate the renderer\n", + "renderer = mujoco.Renderer(env.model)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "9f2ME2WbA5Ip" + }, + "outputs": [], + "source": [ + "#@title Define a render utility function\n", + "\n", + "def get_image(state: State, camera: str) -\u003e np.ndarray:\n", + " \"\"\"Renders the environment state.\"\"\"\n", + " d = mujoco.MjData(env.model)\n", + " # write the mjx.Data into an mjData object\n", + " mjx.device_get_into(d, state.pipeline_state)\n", + " mujoco.mj_forward(env.model, d)\n", + " # use the mjData object to update the renderer\n", + " renderer.update_scene(d, camera=camera)\n", + " return renderer.render()\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "Ph8u-v2Q2xLS" + }, + "outputs": [], + "source": [ + "# initialize the state\n", + "state = jit_reset(jax.random.PRNGKey(0))\n", + "rollout = [state]\n", + "images = [get_image(state, camera='side')]\n", + "\n", + "# grab a trajectory\n", + "for i in range(10):\n", + " ctrl = -0.1 * jp.ones(env.sys.nu)\n", + " state = jit_step(state, ctrl)\n", + " rollout.append(state)\n", + " images.append(get_image(state, camera='side'))\n", + "\n", + "media.show_video(images, fps=1.0 / env.dt)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "BQDG6NQ1CbZD" + }, + "source": [ + "## Train Humanoid Policy\n", + "\n", + "Let's finally train a policy with PPO to make the Humanoid run forwards. Training takes about 13-14 minutes on a Tesla V100 GPU." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "xLiddQYPApBw" + }, + "outputs": [], + "source": [ + "train_fn = functools.partial(\n", + " ppo.train, num_timesteps=30_000_000, num_evals=5, reward_scaling=0.1,\n", + " episode_length=1000, normalize_observations=True, action_repeat=1,\n", + " unroll_length=10, num_minibatches=32, num_updates_per_batch=8,\n", + " discounting=0.97, learning_rate=3e-4, entropy_cost=1e-3, num_envs=2048,\n", + " batch_size=1024, seed=0)\n", + "\n", + "\n", + "x_data = []\n", + "y_data = []\n", + "ydataerr = []\n", + "times = [datetime.now()]\n", + "\n", + "max_y, min_y = 13000, 0\n", + "def progress(num_steps, metrics):\n", + " times.append(datetime.now())\n", + " x_data.append(num_steps)\n", + " y_data.append(metrics['eval/episode_reward'])\n", + " ydataerr.append(metrics['eval/episode_reward_std'])\n", + "\n", + " plt.xlim([0, train_fn.keywords['num_timesteps'] * 1.25])\n", + " plt.ylim([min_y, max_y])\n", + "\n", + " plt.xlabel('# environment steps')\n", + " plt.ylabel('reward per episode')\n", + " plt.title(f'y={y_data[-1]:.3f}')\n", + "\n", + " plt.errorbar(\n", + " x_data, y_data, yerr=ydataerr)\n", + " plt.show()\n", + "\n", + "make_inference_fn, params, _= train_fn(environment=env, progress_fn=progress)\n", + "\n", + "print(f'time to jit: {times[1] - times[0]}')\n", + "print(f'time to train: {times[-1] - times[1]}')" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "YYIch0HEApBx" + }, + "source": [ + "## Save and Load Policy\n", + "\n", + "We can save and load the policy using the brax model API." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "Z8gI6qH6ApBx" + }, + "outputs": [], + "source": [ + "#@title Save Model\n", + "model_path = '/tmp/mjx_brax_policy'\n", + "model.save_params(model_path, params)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "h4reaWgxApBx" + }, + "outputs": [], + "source": [ + "#@title Load Model and Define Inference Function\n", + "params = model.load_params(model_path)\n", + "\n", + "inference_fn = make_inference_fn(params)\n", + "jit_inference_fn = jax.jit(inference_fn)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "0G357XIfApBy" + }, + "source": [ + "## Visualize Policy\n", + "\n", + "Finally we can visualize the policy." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "osYasMw4ApBy" + }, + "outputs": [], + "source": [ + "eval_env = envs.get_environment(env_name)\n", + "\n", + "jit_reset = jax.jit(eval_env.reset)\n", + "jit_step = jax.jit(eval_env.step)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "d-UhypudApBy" + }, + "outputs": [], + "source": [ + "# initialize the state\n", + "rng = jax.random.PRNGKey(0)\n", + "state = jit_reset(rng)\n", + "rollout = [state]\n", + "images = [get_image(state, camera='side')]\n", + "\n", + "# grab a trajectory\n", + "n_steps = 500\n", + "render_every = 2\n", + "\n", + "for i in range(n_steps):\n", + " act_rng, rng = jax.random.split(rng)\n", + " ctrl, _ = jit_inference_fn(state.obs, act_rng)\n", + " state = jit_step(state, ctrl)\n", + " rollout.append(state)\n", + " if i % render_every == 0:\n", + " images.append(get_image(state, camera='side'))\n", + "\n", + " if state.done:\n", + " break\n", + "\n", + "media.show_video(images, fps=1.0 / eval_env.dt / render_every)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "zR-heox6LARK" + }, + "source": [ + "# MJX Policy in MuJoCo\n", + "\n", + "Note that we can also perform the physics step using the original MuJoCo python bindings to show that the policy trained in MJX works in MuJoCo." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "w6ixFi4dApBy" + }, + "outputs": [], + "source": [ + "mj_model = eval_env.model\n", + "mj_data = mujoco.MjData(mj_model)\n", + "\n", + "renderer = mujoco.Renderer(mj_model)\n", + "ctrl = jp.zeros(mj_model.nu)\n", + "\n", + "images = []\n", + "for i in range(n_steps):\n", + " act_rng, rng = jax.random.split(rng)\n", + "\n", + " obs = eval_env._get_obs(mjx.device_put(mj_data), ctrl)\n", + " ctrl, _ = jit_inference_fn(obs, act_rng)\n", + "\n", + " mj_data.ctrl = ctrl\n", + " for _ in range(eval_env._physics_steps_per_control_step):\n", + " mujoco.mj_step(mj_model, mj_data) # Physics step using MuJoCo mj_step.\n", + "\n", + " if i % render_every == 0:\n", + " renderer.update_scene(mj_data, camera='side')\n", + " images.append(renderer.render())\n", + "\n", + "media.show_video(images, fps=1.0 / eval_env.dt / render_every)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "65mIPj6DQNNa" + }, + "source": [ + "# Domain Randomization\n", + "\n", + "We might also want to include randomization over certain `mjModel` parameters while training a policy. In MJX, we can easily create a batch of environments with randomized values populated in `mjx.Model`. Below, we show a function that randomizes friction and actuator gain/bias." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "h8mhzKjHQuoL" + }, + "outputs": [], + "source": [ + "def domain_randomize(sys, rng):\n", + " \"\"\"Randomizes the mjx.Model.\"\"\"\n", + " @jax.vmap\n", + " def rand(rng):\n", + " _, key = jax.random.split(rng, 2)\n", + " # friction\n", + " friction = jax.random.uniform(key, (1,), minval=0.6, maxval=1.4)\n", + " friction = sys.geom_friction.at[:, 0].set(friction)\n", + " # actuator\n", + " _, key = jax.random.split(key, 2)\n", + " gain_range = (-10, -5)\n", + " param = jax.random.uniform(\n", + " key, (1,), minval=gain_range[0], maxval=gain_range[1]\n", + " ) + sys.actuator_gainprm[:, 0]\n", + " gain = sys.actuator_gainprm.at[:, 0].set(param)\n", + " bias = sys.actuator_biasprm.at[:, 1].set(-param)\n", + " return friction, gain, bias\n", + "\n", + " friction, gain, bias = rand(rng)\n", + "\n", + " in_axes = jax.tree_map(lambda x: None, sys)\n", + " in_axes = in_axes.tree_replace({\n", + " 'geom_friction': 0,\n", + " 'actuator_gainprm': 0,\n", + " 'actuator_biasprm': 0,\n", + " })\n", + "\n", + " sys = sys.tree_replace({\n", + " 'geom_friction': friction,\n", + " 'actuator_gainprm': gain,\n", + " 'actuator_biasprm': bias,\n", + " })\n", + "\n", + " return sys, in_axes" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "gnsZo-GWSYYj" + }, + "source": [ + "If we wanted 10 environments with randomized friction and actuator params, we can call `domain_randomize`, which returns a batched `mjModel` along with a dictionary specifying the axes that are batched." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "1K45Kp2ASV9s" + }, + "outputs": [], + "source": [ + "rng = jax.random.PRNGKey(0)\n", + "rng = jax.random.split(rng, 10)\n", + "batched_sys, _ = domain_randomize(env.sys, rng)\n", + "\n", + "print('Single env friction shape: ', env.sys.geom_friction.shape)\n", + "print('Batched env friction shape: ', batched_sys.geom_friction.shape)\n", + "\n", + "print('Friction on geom 0: ', env.sys.geom_friction[0, 0])\n", + "print('Random frictions on geom 0: ', batched_sys.geom_friction[:, 0, 0])" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "efnxNOnpQFuC" + }, + "source": [ + "## Quadruped Env\n", + "\n", + "Let's define a quadruped environment that takes advantage of the domain randomization function. Here we use the [Barkour v0 Quadruped](https://github.com/google-deepmind/mujoco_menagerie/tree/main/google_barkour_v0) and an environment that trains a joystick policy." + ] + }, { "cell_type": "code", "execution_count": null, @@ -483,13 +869,14 @@ " # Penalize non-zero roll and pitch angles. L2 penalty.\n", " orientation=-5.0,\n", " # L2 regularization of joint torques, |tau|^2.\n", - " torques=-0.0002,\n", + " # torques=-0.0002,\n", + " torques=-0.002,\n", " # Penalize the change in the action and encourage smooth\n", " # actions. L2 regularization |action - last_action|^2\n", - " action_rate=-0.3,\n", + " action_rate=-0.1,\n", " # Encourage long swing steps. However, it does not\n", " # encourage high clearances.\n", - " feet_air_time=0.1,\n", + " feet_air_time=0.2,\n", " # Encourage no motion at zero command, L2 regularization\n", " # |q - q_default|^2.\n", " stand_still=-0.5,\n", @@ -550,7 +937,7 @@ " self.uppers = self._default_ap_pose + jp.array([0.2, 0.8, 0.8] * 4)\n", "\n", " def sample_command(self, rng: jax.Array) -\u003e jax.Array:\n", - " lin_vel_x = [-0.6, 1.5] # min max [m/s]\n", + " lin_vel_x = [-0.6, 1.0] # min max [m/s]\n", " lin_vel_y = [-0.8, 0.8] # min max [m/s]\n", " ang_vel_yaw = [-0.7, 0.7] # min max [rad/s]\n", "\n", @@ -857,171 +1244,19 @@ "envs.register_environment('barkour', BarkourEnv)" ] }, - { - "cell_type": "markdown", - "metadata": { - "id": "P1K6IznI2y83" - }, - "source": [ - "## Visualize a Rollout\n", - "\n", - "Let's visualize an environment rollout by running a few env steps with sinusoidal actuation." - ] - }, { "cell_type": "code", "execution_count": null, "metadata": { - "cellView": "form", - "id": "EhKLFK54C1CH" + "id": "pi_yrcz-Qp3W" }, "outputs": [], "source": [ - "# instantiate the environment\n", - "\n", - "env_name = 'humanoid' # @param ['barkour', 'humanoid']\n", + "env_name = 'barkour'\n", "env = envs.get_environment(env_name)\n", "\n", - "# define the jit reset/step functions\n", - "jit_reset = jax.jit(env.reset)\n", - "jit_step = jax.jit(env.step)\n", - "\n", - "# instantiate the renderer\n", - "renderer = mujoco.Renderer(env.model)\n", - "\n", - "def get_image(state: State):\n", - " d = mujoco.MjData(env.model)\n", - " # write the mjx.Data into an mjData object\n", - " mjx.device_get_into(d, state.pipeline_state)\n", - " mujoco.mj_forward(env.model, d)\n", - " if env_name == 'barkour':\n", - " camera='track'\n", - " elif env_name == 'humanoid':\n", - " camera = 'side'\n", - " else:\n", - " raise NotImplementedError(env_name)\n", - " # use the mjData object to update the renderer\n", - " renderer.update_scene(d, camera=camera)\n", - " return renderer.render()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "Ph8u-v2Q2xLS" - }, - "outputs": [], - "source": [ - "# initialize the state\n", - "state = jit_reset(jax.random.PRNGKey(0))\n", - "rollout = [state]\n", - "images = [get_image(state)]\n", - "\n", - "# grab a trajectory\n", - "n_steps = 500\n", - "render_every = 5\n", - "\n", - "for i in range(n_steps):\n", - " ctrl = jp.ones(env.sys.nu) * jp.sin(2.0 * jp.pi * i / 500)\n", - " state = jit_step(state, ctrl)\n", - " rollout.append(state)\n", - " if i % render_every == 0:\n", - " images.append(get_image(state))\n", - "\n", - "media.show_video(images, fps=1.0 / env.dt / render_every)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "65mIPj6DQNNa" - }, - "source": [ - "# Domain Randomization\n", - "\n", - "We have all the pieces to train a policy, but we might also want to include randomization over certain `mjModel` parameters. In MJX, we can easily create a batch of environments with randomized values populated in `mjx.Model`. Below, we show a function that randomizes friction and actuator gain/bias." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "h8mhzKjHQuoL" - }, - "outputs": [], - "source": [ - "def domain_randomize(sys, rng):\n", - " \"\"\"Randomizes the mjx.Model.\"\"\"\n", - " @jax.vmap\n", - " def rand(rng):\n", - " _, key = jax.random.split(rng, 2)\n", - " # friction\n", - " friction = jax.random.uniform(key, (1,), minval=0.8, maxval=1.2)\n", - " friction = sys.geom_friction.at[:, 0].set(friction)\n", - " # actuator\n", - " _, key = jax.random.split(key, 2)\n", - " gain_range = (-5, 5)\n", - " param = jax.random.uniform(\n", - " key, (1,), minval=gain_range[0], maxval=gain_range[1]\n", - " ) + sys.actuator_gainprm[:, 0]\n", - " gain = sys.actuator_gainprm.at[:, 0].set(param)\n", - " bias = sys.actuator_biasprm.at[:, 1].set(-param)\n", - " return friction, gain, bias\n", - "\n", - " friction, gain, bias = rand(rng)\n", - "\n", - " in_axes = jax.tree_map(lambda x: None, sys)\n", - " in_axes = in_axes.tree_replace({\n", - " 'geom_friction': 0,\n", - " 'actuator_gainprm': 0,\n", - " 'actuator_biasprm': 0,\n", - " })\n", - "\n", - " sys = sys.tree_replace({\n", - " 'geom_friction': friction,\n", - " 'actuator_gainprm': gain,\n", - " 'actuator_biasprm': bias,\n", - " })\n", - "\n", - " return sys, in_axes" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "gnsZo-GWSYYj" - }, - "source": [ - "If we wanted 10 environments with randomized friction and actuator params, we can call `domain_randomize`." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "1K45Kp2ASV9s" - }, - "outputs": [], - "source": [ - "rng = jax.random.PRNGKey(0)\n", - "rng = jax.random.split(rng, 10)\n", - "batched_sys, _ = domain_randomize(env.sys, rng)\n", - "\n", - "print('Single env friction shape: ', env.sys.geom_friction.shape)\n", - "print('Batched env friction shape: ', batched_sys.geom_friction.shape)\n", - "\n", - "print('Friction on geom 0: ', env.sys.geom_friction[0, 0])\n", - "print('Random frictions on geom 0: ', batched_sys.geom_friction[:, 0, 0])" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "LdqPYBlPVWwc" - }, - "source": [ - "Note that we'll pass the `domain_randomize` function into the brax trainer; the brax training algorithms will call the randomizer over the environment batch dimension." + "# re-instantiate the renderer\n", + "renderer = mujoco.Renderer(env.model)" ] }, { @@ -1030,9 +1265,9 @@ "id": "nxaNFP9mA23H" }, "source": [ - "# Train Policy\n", + "## Train Policy\n", "\n", - "Let's finally train a policy with PPO. Training takes about 12-13 minutes on a Tesla V100 GPU for both environments." + "To train a policy with domain randomization, we pass in the domain randomization function into the brax train function; brax will call the domain randomization function when rolling out episodes. Training the quadruped takes about 14 minutes on a Tesla V100 GPU." ] }, { @@ -1043,59 +1278,28 @@ }, "outputs": [], "source": [ + "make_networks_factory = functools.partial(\n", + " ppo_networks.make_ppo_networks,\n", + " policy_hidden_layer_sizes=(128, 128, 128, 128))\n", + "train_fn = functools.partial(\n", + " ppo.train,\n", + " num_timesteps=60_000_000, num_evals=3, reward_scaling=1,\n", + " episode_length=1000, normalize_observations=True,\n", + " action_repeat=1, unroll_length=20, num_minibatches=8, gae_lambda=0.95,\n", + " num_updates_per_batch=4, discounting=0.99, learning_rate=3e-4,\n", + " entropy_cost=1e-2, num_envs=8192, batch_size=1024,\n", + " network_factory=make_networks_factory,\n", + " num_resets_per_eval=10,\n", + " randomization_fn=domain_randomize, seed=0)\n", "\n", - "if env_name == 'barkour':\n", - " make_networks_factory = functools.partial(\n", - " ppo_networks.make_ppo_networks,\n", - " policy_hidden_layer_sizes=(128, 128, 128, 128))\n", - " train_fn = functools.partial(\n", - " ppo.train,\n", - " num_timesteps=60_000_000, num_evals=3, reward_scaling=1,\n", - " episode_length=1000, normalize_observations=True,\n", - " action_repeat=1, unroll_length=20, num_minibatches=32, gae_lambda=0.95,\n", - " num_updates_per_batch=4, discounting=0.99, learning_rate=3.0e-4,\n", - " entropy_cost=1e-2, num_envs=8192, batch_size=1024,\n", - " network_factory=make_networks_factory,\n", - " num_resets_per_eval=10,\n", - " randomization_fn=domain_randomize, seed=0)\n", - "elif env_name == 'humanoid':\n", - " train_fn = functools.partial(\n", - " ppo.train, num_timesteps=30_000_000, num_evals=5, reward_scaling=0.1,\n", - " episode_length=1000, normalize_observations=True, action_repeat=1,\n", - " unroll_length=10, num_minibatches=32, num_updates_per_batch=8,\n", - " discounting=0.97, learning_rate=3e-4, entropy_cost=1e-3, num_envs=2048,\n", - " batch_size=1024, seed=0)\n", - "else:\n", - " raise NotImplementedError(\n", - " f'env_name: {env_name} is not implemented in this notebook.')\n", - "\n", - "\n", - "max_y = {'barkour':30, 'humanoid': 13000}[env_name]\n", - "min_y = {'barkour': -15}.get(env_name, 0)\n", "\n", "x_data = []\n", "y_data = []\n", "ydataerr = []\n", "times = [datetime.now()]\n", + "max_y, min_y = 30, 0\n", "\n", - "def progress(num_steps, metrics):\n", - " times.append(datetime.now())\n", - " x_data.append(num_steps)\n", - " y_data.append(metrics['eval/episode_reward'])\n", - " ydataerr.append(metrics['eval/episode_reward_std'])\n", - "\n", - " plt.xlim([0, train_fn.keywords['num_timesteps'] * 1.25])\n", - " plt.ylim([min_y, max_y])\n", - "\n", - " plt.xlabel('# environment steps')\n", - " plt.ylabel('reward per episode')\n", - " plt.title(f'y={y_data[-1]:.3f}')\n", - "\n", - " plt.errorbar(\n", - " x_data, y_data, yerr=ydataerr)\n", - " plt.show()\n", - "\n", - "# Reset environments since internals may be overwritten by tracers due to\n", + "# Reset environments since internals may be overwritten by tracers from the\n", "# domain randomization function.\n", "env = envs.get_environment(env_name)\n", "eval_env = envs.get_environment(env_name)\n", @@ -1107,30 +1311,6 @@ "print(f'time to train: {times[-1] - times[1]}')" ] }, - { - "cell_type": "markdown", - "metadata": { - "id": "-dRmWHsHB_1K" - }, - "source": [ - "# Save and Load Policy\n", - "\n", - "We can save and load the policy using the brax model API." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "u9sdN4Xa5JoH" - }, - "outputs": [], - "source": [ - "#@title Save Model\n", - "model_path = '/tmp/mjx_brax_policy'\n", - "model.save_params(model_path, params)" - ] - }, { "cell_type": "code", "execution_count": null, @@ -1139,7 +1319,9 @@ }, "outputs": [], "source": [ - "#@title Load Model and Define Inference Function\n", + "# Save and reload params.\n", + "model_path = '/tmp/mjx_brax_quadruped_policy'\n", + "model.save_params(model_path, params)\n", "params = model.load_params(model_path)\n", "\n", "inference_fn = make_inference_fn(params)\n", @@ -1152,9 +1334,9 @@ "id": "L01IrN4oCIkC" }, "source": [ - "# Visualize Policy\n", + "## Visualize Policy\n", "\n", - "Finally we can visualize the policy. For the Barkour Quadruped, the joystick commands can be set through `x_vel`, `y_vel`, and `ang_vel`. `x_vel` and `y_vel` define the linear forward and sideways velocities with respect to the quadruped torso. `ang_vel` defines the angular velocity of the torso in the z direction." + "For the Barkour Quadruped, the joystick commands can be set through `x_vel`, `y_vel`, and `ang_vel`. `x_vel` and `y_vel` define the linear forward and sideways velocities with respect to the quadruped torso. `ang_vel` defines the angular velocity of the torso in the z direction." ] }, { @@ -1183,7 +1365,7 @@ "# @markdown Commands **only used for Barkour Env**:\n", "x_vel = 1.0 #@param {type: \"number\"}\n", "y_vel = 0.0 #@param {type: \"number\"}\n", - "ang_vel = -0.0 #@param {type: \"number\"}\n", + "ang_vel = -0.5 #@param {type: \"number\"}\n", "\n", "the_command = jp.array([x_vel, y_vel, ang_vel])\n", "\n", @@ -1192,7 +1374,7 @@ "state = jit_reset(rng)\n", "state.info['command'] = the_command\n", "rollout = [state]\n", - "images = [get_image(state)]\n", + "images = [get_image(state, camera='track')]\n", "\n", "# grab a trajectory\n", "n_steps = 500\n", @@ -1204,7 +1386,7 @@ " state = jit_step(state, ctrl)\n", " rollout.append(state)\n", " if i % render_every == 0:\n", - " images.append(get_image(state))\n", + " images.append(get_image(state, camera='track'))\n", "\n", "media.show_video(images, fps=1.0 / eval_env.dt / render_every)" ] @@ -1216,7 +1398,12 @@ "gpuClass": "premium", "gpuType": "V100", "private_outputs": true, - "provenance": [], + "provenance": [ + { + "file_id": "1brcF4_qCRS2ASc-QQw1rsEwl5IjzGvq2", + "timestamp": 1697763780236 + } + ], "toc_visible": true }, "kernelspec": { From f6ed57e9e8f84f080eaff3755d361a3c09c4264b Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Sun, 22 Oct 2023 08:18:38 -0700 Subject: [PATCH 11/74] Fix typos in changelog PiperOrigin-RevId: 575618675 Change-Id: I51a8e5aee124c632ff17d0cb0610ca95822bf859 --- doc/changelog.rst | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index e64a7300..859d1558 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -8,14 +8,10 @@ Upcoming version (not yet released) Bug fixes ^^^^^^^^^ -1. Fix in simulate: correct handling of "Pause update", "Fullscreen" and "VSync" buttons. - -MJX Bug fixes -^^^^^^^^^ - -1. Fix typos and supported fields in the docs (fixes :github:issue:`1105` and :github:issue:`1106`). -2. Fix bug where mixed `jnt_limited` joints are not being constrained correctly. -3. Make `device_put` type validation more verbose (fixed :github:issue:`1113`). +1. :ref:`simulate`: correct handling of "Pause update", "Fullscreen" and "VSync" buttons. +2. Fixed typos and supported fields in docs (fixes :github:issue:`1105` and :github:issue:`1106`). +3. Fixed bug where mixed ``jnt_limited`` joints were not being constrained correctly. +4. Made ``device_put`` type validation more verbose (fixes :github:issue:`1113`). Version 3.0.0 (October 18, 2023) -------------------------------- From 179ea8ef03bb63b15035f4bf69341a19f772ef23 Mon Sep 17 00:00:00 2001 From: Saran Tunyasuvunakool Date: Sun, 22 Oct 2023 08:19:18 -0700 Subject: [PATCH 12/74] Add MJX to the doc title. This should make the listing clearer on search engines. PiperOrigin-RevId: 575618744 Change-Id: I79feea95c31466dfb373e42a547116044ce2ac08 --- doc/mjx.rst | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/doc/mjx.rst b/doc/mjx.rst index 92566ddb..0ad90ffa 100644 --- a/doc/mjx.rst +++ b/doc/mjx.rst @@ -1,6 +1,6 @@ -========== -MuJoCo XLA -========== +================ +MuJoCo XLA (MJX) +================ Starting with version 3.0.0, MuJoCo includes MuJoCo XLA (MJX) under the `mjx `__ directory. MJX allows MuJoCo to run on compute From a53b5be956f6224784fb84b4149abb775ef8d61b Mon Sep 17 00:00:00 2001 From: Kevin Zakka Date: Sun, 22 Oct 2023 10:18:36 -0700 Subject: [PATCH 13/74] Allow pre-release `mujoco` installs in `mjx` source install. PiperOrigin-RevId: 575629464 Change-Id: I47615dee0e965c6016d90871fa100f52bb74bf0c --- mjx/pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mjx/pyproject.toml b/mjx/pyproject.toml index d65fab25..bea4b310 100644 --- a/mjx/pyproject.toml +++ b/mjx/pyproject.toml @@ -31,7 +31,7 @@ dependencies = [ "etils[epath]", "jax", "jaxlib", - "mujoco>=3.0.1", + "mujoco>=3.0.1.dev0", "scipy", "trimesh", ] From c6ef4f71c3ee702e94e206f982766c443f600b8d Mon Sep 17 00:00:00 2001 From: Alessio Quaglino Date: Mon, 23 Oct 2023 04:43:53 -0700 Subject: [PATCH 14/74] Add flex fields to mj_validateReferences. PiperOrigin-RevId: 575779915 Change-Id: I14d22092e7d712dcc3d474cf0c2a85ef4d6a4aee --- src/engine/engine_io.c | 157 ++++++++++++++++++++++------------------- 1 file changed, 84 insertions(+), 73 deletions(-) diff --git a/src/engine/engine_io.c b/src/engine/engine_io.c index 773a20f4..6c0c3198 100644 --- a/src/engine/engine_io.c +++ b/src/engine/engine_io.c @@ -1829,79 +1829,90 @@ const char* mj_validateReferences(const mjModel* m) { // add flex fields (b/303056369) -#define MJMODEL_REFERENCES \ - X(body_parentid, nbody, nbody , 0 ) \ - X(body_rootid, nbody, nbody , 0 ) \ - X(body_weldid, nbody, nbody , 0 ) \ - X(body_mocapid, nbody, nmocap , 0 ) \ - X(body_jntadr, nbody, njnt , m->body_jntnum ) \ - X(body_dofadr, nbody, nv , m->body_dofnum ) \ - X(body_geomadr, nbody, ngeom , m->body_geomnum ) \ - X(body_bvhadr, nbody, nbvh , m->body_bvhnum ) \ - X(body_plugin, nbody, nplugin , 0 ) \ - X(jnt_qposadr, njnt, nq , 0 ) \ - X(jnt_dofadr, njnt, nv , 0 ) \ - X(jnt_bodyid, njnt, nbody , 0 ) \ - X(dof_bodyid, nv, nbody , 0 ) \ - X(dof_jntid, nv, njnt , 0 ) \ - X(dof_parentid, nv, nv , 0 ) \ - X(dof_Madr, nv, nM , 0 ) \ - X(geom_bodyid, ngeom, nbody , 0 ) \ - X(geom_matid, ngeom, nmat , 0 ) \ - X(site_bodyid, nsite, nbody , 0 ) \ - X(site_matid, nsite, nmat , 0 ) \ - X(cam_bodyid, ncam, nbody , 0 ) \ - X(cam_targetbodyid, ncam, nbody , 0 ) \ - X(light_bodyid, nlight, nbody , 0 ) \ - X(light_targetbodyid, nlight, nbody , 0 ) \ - X(mesh_vertadr, nmesh, nmeshvert , m->mesh_vertnum ) \ - X(mesh_normaladr, nmesh, nmeshnormal , m->mesh_normalnum ) \ - X(mesh_texcoordadr, nmesh, nmeshtexcoord, m->mesh_texcoordnum ) \ - X(mesh_faceadr, nmesh, nmeshface , m->mesh_facenum ) \ - X(mesh_bvhadr, nmesh, nbvh , m->mesh_bvhnum ) \ - X(mesh_graphadr, nmesh, nmeshgraph , 0 ) \ - X(skin_matid, nskin, nmat , 0 ) \ - X(skin_vertadr, nskin, nskinvert , m->skin_vertnum ) \ - X(skin_texcoordadr, nskin, nskintexvert , 0 ) \ - X(skin_faceadr, nskin, nskinface , m->skin_facenum ) \ - X(skin_boneadr, nskin, nskinbone , m->skin_bonenum ) \ - X(skin_bonevertadr, nskinbone, nskinbonevert, m->skin_bonevertnum ) \ - X(skin_bonebodyid, nskinbone, nbody , 0 ) \ - X(skin_bonevertid, nskinbonevert, nskinvert , 0 ) \ - X(pair_geom1, npair, ngeom , 0 ) \ - X(pair_geom2, npair, ngeom , 0 ) \ - X(actuator_plugin, nu, nplugin , 0 ) \ - X(actuator_actadr, nu, na , m->actuator_actnum ) \ - X(sensor_plugin, nsensor, nplugin , 0 ) \ - X(plugin_stateadr, nplugin, npluginstate , m->plugin_statenum ) \ - X(plugin_attradr, nplugin, npluginattr , 0 ) \ - X(tendon_adr, ntendon, nwrap , m->tendon_num ) \ - X(tendon_matid, ntendon, nmat , 0 ) \ - X(numeric_adr, nnumeric, nnumericdata , m->numeric_size ) \ - X(text_adr, ntext, ntextdata , m->text_size ) \ - X(tuple_adr, ntuple, ntupledata , m->tuple_size ) \ - X(name_bodyadr, nbody, nnames , 0 ) \ - X(name_jntadr, njnt, nnames , 0 ) \ - X(name_geomadr, ngeom, nnames , 0 ) \ - X(name_siteadr, nsite, nnames , 0 ) \ - X(name_camadr, ncam, nnames , 0 ) \ - X(name_lightadr, nlight, nnames , 0 ) \ - X(name_meshadr, nmesh, nnames , 0 ) \ - X(name_skinadr, nskin, nnames , 0 ) \ - X(name_hfieldadr, nhfield, nnames , 0 ) \ - X(name_texadr, ntex, nnames , 0 ) \ - X(name_matadr, nmat, nnames , 0 ) \ - X(name_pairadr, npair, nnames , 0 ) \ - X(name_excludeadr, nexclude, nnames , 0 ) \ - X(name_eqadr, neq, nnames , 0 ) \ - X(name_tendonadr, ntendon, nnames , 0 ) \ - X(name_actuatoradr, nu, nnames , 0 ) \ - X(name_sensoradr, nsensor, nnames , 0 ) \ - X(name_numericadr, nnumeric, nnames , 0 ) \ - X(name_textadr, ntext, nnames , 0 ) \ - X(name_tupleadr, ntuple, nnames , 0 ) \ - X(name_keyadr, nkey, nnames , 0 ) \ - X(mesh_pathadr, nmesh, npaths , 0 ) +#define MJMODEL_REFERENCES \ + X(body_parentid, nbody, nbody , 0 ) \ + X(body_rootid, nbody, nbody , 0 ) \ + X(body_weldid, nbody, nbody , 0 ) \ + X(body_mocapid, nbody, nmocap , 0 ) \ + X(body_jntadr, nbody, njnt , m->body_jntnum ) \ + X(body_dofadr, nbody, nv , m->body_dofnum ) \ + X(body_geomadr, nbody, ngeom , m->body_geomnum ) \ + X(body_bvhadr, nbody, nbvh , m->body_bvhnum ) \ + X(body_plugin, nbody, nplugin , 0 ) \ + X(jnt_qposadr, njnt, nq , 0 ) \ + X(jnt_dofadr, njnt, nv , 0 ) \ + X(jnt_bodyid, njnt, nbody , 0 ) \ + X(dof_bodyid, nv, nbody , 0 ) \ + X(dof_jntid, nv, njnt , 0 ) \ + X(dof_parentid, nv, nv , 0 ) \ + X(dof_Madr, nv, nM , 0 ) \ + X(geom_bodyid, ngeom, nbody , 0 ) \ + X(geom_matid, ngeom, nmat , 0 ) \ + X(site_bodyid, nsite, nbody , 0 ) \ + X(site_matid, nsite, nmat , 0 ) \ + X(cam_bodyid, ncam, nbody , 0 ) \ + X(cam_targetbodyid, ncam, nbody , 0 ) \ + X(light_bodyid, nlight, nbody , 0 ) \ + X(light_targetbodyid, nlight, nbody , 0 ) \ + X(mesh_vertadr, nmesh, nmeshvert , m->mesh_vertnum ) \ + X(mesh_normaladr, nmesh, nmeshnormal , m->mesh_normalnum ) \ + X(mesh_texcoordadr, nmesh, nmeshtexcoord , m->mesh_texcoordnum ) \ + X(mesh_faceadr, nmesh, nmeshface , m->mesh_facenum ) \ + X(mesh_bvhadr, nmesh, nbvh , m->mesh_bvhnum ) \ + X(mesh_graphadr, nmesh, nmeshgraph , 0 ) \ + X(flex_vertadr, nflex, nflexvert , m->flex_vertnum ) \ + X(flex_edgeadr, nflex, nflexedge , m->flex_edgenum ) \ + X(flex_elemadr, nflex, nflexelem , m->flex_elemnum ) \ + X(flex_evpairadr, nflex, nflexevpair , m->flex_evpairnum ) \ + X(flex_texcoordadr, nflex, nflextexcoord , 0 ) \ + X(flex_elemdataadr, nflex, nflexelemdata , 0 ) \ + X(flex_shelldataadr, nflex, nflexshelldata, 0 ) \ + X(flex_edge, nflexedge*2, nflexvert , 0 ) \ + X(flex_elem, nflexelemdata, nflexvert , 0 ) \ + X(flex_shell, nflexshelldata, nflexvert , 0 ) \ + X(flex_bvhadr, nflex, nbvh , m->flex_bvhnum ) \ + X(skin_matid, nskin, nmat , 0 ) \ + X(skin_vertadr, nskin, nskinvert , m->skin_vertnum ) \ + X(skin_texcoordadr, nskin, nskintexvert , 0 ) \ + X(skin_faceadr, nskin, nskinface , m->skin_facenum ) \ + X(skin_boneadr, nskin, nskinbone , m->skin_bonenum ) \ + X(skin_bonevertadr, nskinbone, nskinbonevert , m->skin_bonevertnum ) \ + X(skin_bonebodyid, nskinbone, nbody , 0 ) \ + X(skin_bonevertid, nskinbonevert, nskinvert , 0 ) \ + X(pair_geom1, npair, ngeom , 0 ) \ + X(pair_geom2, npair, ngeom , 0 ) \ + X(actuator_plugin, nu, nplugin , 0 ) \ + X(actuator_actadr, nu, na , m->actuator_actnum ) \ + X(sensor_plugin, nsensor, nplugin , 0 ) \ + X(plugin_stateadr, nplugin, npluginstate , m->plugin_statenum ) \ + X(plugin_attradr, nplugin, npluginattr , 0 ) \ + X(tendon_adr, ntendon, nwrap , m->tendon_num ) \ + X(tendon_matid, ntendon, nmat , 0 ) \ + X(numeric_adr, nnumeric, nnumericdata , m->numeric_size ) \ + X(text_adr, ntext, ntextdata , m->text_size ) \ + X(tuple_adr, ntuple, ntupledata , m->tuple_size ) \ + X(name_bodyadr, nbody, nnames , 0 ) \ + X(name_jntadr, njnt, nnames , 0 ) \ + X(name_geomadr, ngeom, nnames , 0 ) \ + X(name_siteadr, nsite, nnames , 0 ) \ + X(name_camadr, ncam, nnames , 0 ) \ + X(name_lightadr, nlight, nnames , 0 ) \ + X(name_meshadr, nmesh, nnames , 0 ) \ + X(name_skinadr, nskin, nnames , 0 ) \ + X(name_hfieldadr, nhfield, nnames , 0 ) \ + X(name_texadr, ntex, nnames , 0 ) \ + X(name_matadr, nmat, nnames , 0 ) \ + X(name_pairadr, npair, nnames , 0 ) \ + X(name_excludeadr, nexclude, nnames , 0 ) \ + X(name_eqadr, neq, nnames , 0 ) \ + X(name_tendonadr, ntendon, nnames , 0 ) \ + X(name_actuatoradr, nu, nnames , 0 ) \ + X(name_sensoradr, nsensor, nnames , 0 ) \ + X(name_numericadr, nnumeric, nnames , 0 ) \ + X(name_textadr, ntext, nnames , 0 ) \ + X(name_tupleadr, ntuple, nnames , 0 ) \ + X(name_keyadr, nkey, nnames , 0 ) \ + X(mesh_pathadr, nmesh, npaths , 0 ) #define X(adrarray, nadrs, ntarget, numarray) { \ int *nums = (numarray); \ From 2bb8652bf6605e63cc7356ee1058e9c93d9654af Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Mon, 23 Oct 2023 06:34:24 -0700 Subject: [PATCH 15/74] Add User Interface documentation. PiperOrigin-RevId: 575800982 Change-Id: I9b2e03932d0d3b1844b90e5da7c32a640d999e69 --- doc/APIreference/APItypes.rst | 18 +++- doc/APIreference/functions.rst | 25 +++++- doc/APIreference/functions_override.rst | 37 +++++++++ doc/changelog.rst | 5 ++ doc/programming/index.rst | 1 + doc/programming/ui.rst | 105 ++++++++++++++++++++++++ 6 files changed, 184 insertions(+), 7 deletions(-) create mode 100644 doc/programming/ui.rst diff --git a/doc/APIreference/APItypes.rst b/doc/APIreference/APItypes.rst index d2a812d6..684e810b 100644 --- a/doc/APIreference/APItypes.rst +++ b/doc/APIreference/APItypes.rst @@ -34,6 +34,8 @@ MuJoCo defines a large number of types: - Structs used by :ref:`engine plugins`. - Several :ref:`tyFunction` for user-defined callbacks. +- :ref:`tyNotes` regarding specific data structures that require detailed description. + .. _tyPrimitive: @@ -909,6 +911,7 @@ This structure contains the custom OpenGL rendering context, with the ids of all User Interface ^^^^^^^^^^^^^^ +For a high-level description of the UI framework, see :ref:`UI`. The names of these struct types are prefixed with ``mjui``, except for the main :ref:`mjUI` struct itself. @@ -917,7 +920,9 @@ The names of these struct types are prefixed with ``mjui``, except for the main mjuiState ~~~~~~~~~ -This structure contains the keyboard and mouse state used by the UI framework. +This C struct represents the global state of the window, keyboard and mouse, input event descriptors, and all window +rectangles (including the visible UI rectangles). There is only one ``mjuiState`` per application, even if there are +multiple UIs. This struct would normally be defined as a global variable. .. mujoco-include:: mjuiState @@ -967,7 +972,9 @@ This structure defines one section of the UI. mjuiDef ~~~~~~~ -This structure defines one entry in the definition table used for simplified UI construction. +This structure defines one entry in the definition table used for simplified UI construction. It contains everything +needed to define one UI item. Some translation is performed by the helper functions, so that multiple mjuiDefs can be +defined as a static table. .. mujoco-include:: mjuiDef @@ -977,7 +984,12 @@ This structure defines one entry in the definition table used for simplified UI mjUI ~~~~ -This structure defines the entire UI. +This C struct represents an entire UI. The same application could have multiple UIs, for example on the left and the +right of the window. This would normally be defined as a global variable. As explained earlier, it contains static +allocation for a maximum number of supported UI sections (:ref:`mjuiSection`) each with a maximum number +of supported items (:ref:`mjuiItem`). It also contains the color and spacing themes, enable/disable +callback, virtual window descriptor, text edit state, mouse focus. Some of these fields are set only once when the UI +is initialized, others change at runtime. .. mujoco-include:: mjUI diff --git a/doc/APIreference/functions.rst b/doc/APIreference/functions.rst index 09622ed0..91c35b1b 100644 --- a/doc/APIreference/functions.rst +++ b/doc/APIreference/functions.rst @@ -1945,6 +1945,8 @@ Find first rectangle containing mouse, -1: not found. UI framework ^^^^^^^^^^^^ +For a high-level description of the UI framework, see :ref:`UI`. + .. _mjui_themeSpacing: mjui_themeSpacing @@ -1970,7 +1972,12 @@ mjui_add .. mujoco-include:: mjui_add -Add definitions to UI. +This is the helper function used to construct a UI. The second argument points to an array of :ref:`mjuiDef` structs, +each corresponding to one item. The last (unused) item has its type set to -1, to mark termination. The items are added +after the end of the last used section. There is also another version of this function +(:ref:`mjui_addToSection`) which adds items to a specified section instead of adding them at the end +of the UI. Keep in mind that there is a maximum preallocated number of sections and items per section, given by +:ref:`mjMAXUISECT` and :ref:`mjMAXUIITEM`. Exceeding these maxima results in low-level errors. .. _mjui_addToSection: @@ -1997,7 +2004,12 @@ mjui_update .. mujoco-include:: mjui_update -Update specific section/item; -1: update all. +This is the main UI update function. It needs to be called whenever the user data (pointed to by the item data pointers) +changes, or when the UI state itself changes. It is normally called by a higher-level function implemented by the user +(``UiModify`` in :ref:`simulate.cc `) which also recomputes the layout of all rectangles and associated +auxiliary buffers. The function updates the pixels in the offscreen OpenGL buffer. To perform minimal updates, the user +specifies the section and the item that was modified. A value of -1 means all items and/or sections need to be updated +(which is needed following major changes.) .. _mjui_event: @@ -2006,7 +2018,10 @@ mjui_event .. mujoco-include:: mjui_event -Handle UI event, return pointer to changed item, NULL if no change. +This function is the low-level event handler. It makes the necessary changes in the UI and returns a pointer to the item +that received the event (or ``NULL`` if no valid event was recorded). This is normally called within the event handler +implemented by the user (``UiEvent`` in :ref:`simulate.cc `), and then some action is taken by user code +depending on which UI item was modified and what the state of that item is after the event is handled. .. _mjui_render: @@ -2015,7 +2030,9 @@ mjui_render .. mujoco-include:: mjui_render -Copy UI image to current buffer. +This function is called in the screen refresh loop. It copies the offscreen OpenGL buffer to the window framebuffer. If +there are multiple UIs in the application, it should be called once for each UI. Thus ``mjui_render`` is called all the +time, while :ref:`mjui_update` is called only when changes in the UI take place. .. _Errorandmemory: diff --git a/doc/APIreference/functions_override.rst b/doc/APIreference/functions_override.rst index 98241ae5..675cb0e2 100644 --- a/doc/APIreference/functions_override.rst +++ b/doc/APIreference/functions_override.rst @@ -218,6 +218,43 @@ of how to use these functions. .. _UIframework: +For a high-level description of the UI framework, see :ref:`UI`. + +.. _mjui_add: + +This is the helper function used to construct a UI. The second argument points to an array of :ref:`mjuiDef` structs, +each corresponding to one item. The last (unused) item has its type set to -1, to mark termination. The items are added +after the end of the last used section. There is also another version of this function +(:ref:`mjui_addToSection`) which adds items to a specified section instead of adding them at the end +of the UI. Keep in mind that there is a maximum preallocated number of sections and items per section, given by +:ref:`mjMAXUISECT` and :ref:`mjMAXUIITEM`. Exceeding these maxima results in low-level errors. + +.. _mjui_update: + +This is the main UI update function. It needs to be called whenever the user data (pointed to by the item data pointers) +changes, or when the UI state itself changes. It is normally called by a higher-level function implemented by the user +(``UiModify`` in :ref:`simulate.cc `) which also recomputes the layout of all rectangles and associated +auxiliary buffers. The function updates the pixels in the offscreen OpenGL buffer. To perform minimal updates, the user +specifies the section and the item that was modified. A value of -1 means all items and/or sections need to be updated +(which is needed following major changes.) + +.. _mjui_event: + +This function is the low-level event handler. It makes the necessary changes in the UI and returns a pointer to the item +that received the event (or ``NULL`` if no valid event was recorded). This is normally called within the event handler +implemented by the user (``UiEvent`` in :ref:`simulate.cc `), and then some action is taken by user code +depending on which UI item was modified and what the state of that item is after the event is handled. + + +.. _mjui_render: + +This function is called in the screen refresh loop. It copies the offscreen OpenGL buffer to the window framebuffer. If +there are multiple UIs in the application, it should be called once for each UI. Thus ``mjui_render`` is called all the +time, while :ref:`mjui_update` is called only when changes in the UI take place. + + + + .. _Errorandmemory: .. _Standardmath: diff --git a/doc/changelog.rst b/doc/changelog.rst index 859d1558..6daa2ae9 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -13,6 +13,11 @@ Bug fixes 3. Fixed bug where mixed ``jnt_limited`` joints were not being constrained correctly. 4. Made ``device_put`` type validation more verbose (fixes :github:issue:`1113`). +Documentation +^^^^^^^^^^^^^ + +5. Added documentation for the :ref:`UI` framework. + Version 3.0.0 (October 18, 2023) -------------------------------- diff --git a/doc/programming/index.rst b/doc/programming/index.rst index b397b80c..381d8754 100644 --- a/doc/programming/index.rst +++ b/doc/programming/index.rst @@ -275,5 +275,6 @@ now lazily resolved at runtime after the switch to GLAD, the "nogl" libraries ar simulation visualization + ui samples extension diff --git a/doc/programming/ui.rst b/doc/programming/ui.rst new file mode 100644 index 00000000..df96ccc3 --- /dev/null +++ b/doc/programming/ui.rst @@ -0,0 +1,105 @@ +.. _UI: + +User Interface +-------------- + +MuJoCo has a native UI framework. Its use is illustrated in the :ref:`simulate.cc ` viewer. It is +designed to be fast in terms of updating and rendering, easy to use both for the developer and for the user, +cross-platform, and integrated with the native MuJoCo renderer. In order to achieve these design goals, we have omitted +many features and customization options that are available in other UI frameworks, and instead focused on efficiency and +automation. + + + +.. _uiDesign: + +Design overview +~~~~~~~~~~~~~~~ + + +Native OpenGL rendering + We do not use any helper tools or libraries. Instead we provide C code for rendering all UI elements directly in + OpenGL. We support multiple UIs, each of which is a virtual rectangle that can be taller than the visible window. The + elements of each UI are rendered offscreen in auxiliary OpenGL buffers, via minimal updates, only when changes are + necessary. At each screen refresh we then copy the pixels from these auxiliary buffers to the window framebuffer, and + also implement vertical scroll bars when the window is smaller than the UI. This copy operation is done on the GPU and + is very fast. + + +Platform abstraction + The software design has 3 layers: OpenGL rendering of UI elements working in conjunction with the MuJoCo renderer + (which is fully cross-platform); abstract functions for access to windows, keyboard and mouse, defined as pure virtual + functions in the ``PlatformUIAdapter`` class; and an implementation of those functions in the derived class + ``GlfwAdapter``. `GLFW `__ itself is cross-platform. Nevertheless, we have opted for this + layered design in order to separate generic from platform-specific functionality. If GLFW needs to be replaced with + another similar framework for some reason, only ``GlfwAdapter`` will need to be rewritten. + + +Themes and appearance + Individual UI elements do not allow customization in terms of appearance or layout. Instead, we use themes for colors + and spacing, and arrange all UI elements automatically. Several built-in themes are provided and the user can design + custom themes, however the entire UI uses a single theme for all elements. Appearance is minimalist: mostly colored + rectangles with text. Bitmaps and other custom decorations are not supported. The UI element types are check boxes, + radio button groups, selection lists, sliders, text edit boxes, static text, buttons, separators. These elements are + grouped into sections which can be expanded and collapsed. + + +Layout and rectangles + Each UI is one virtual rectangle, whose width is determined by the theme and whose height is determined by the + sections, items within each section, and also the expand/collapse state of each section. The sizes and auxiliary + buffers for these virtual rectangles are handled automatically when the UI is updated. Each UI has a visible rectangle + on the screen, and in addition there are other rectangles -- for 3D rendering, 2D figures, and possibly custom OpenGL + rendering. All these visible rectangles are saved (in :ref:`mjuiState`) and are used to determine where mouse events + should be directed. The rectangle layout is updated by a callback provided by the user. + + +Static allocation and creation + Rather than allocating and deallocating a large number of objects corresponding to UI elements and linking them + together, we create a single C struct (type :ref:`mjUI`) with static allocation supporting some maximum number of + sections and elements; and then keep a record of how many are in use. UI creation is simplified by helper functions + whose input is a C struct (type :ref:`mjuiDef`) that is essentially a table where each row describes one UI element + (see below). This makes it possible to construct elaborate user interfaces with surprisingly little C code. + Programmatic UI creation is also possible, for example when populating a UI with sliders corresponding to MuJoCo model + joints. + + +Minimal state + The UI is designed to be as stateless as possible, so as to simplify development. This has two aspects. First, instead + of replicating user data within the UI elements, we store pointers to user data. For example, we might create a UI + slider and set its data pointer to ``mjData* d->qpos+7``. This slider will visualize as well as control the 7th scalar + component of the qpos vector of a MuJoCo model. Thus, when the simulation is updated, we have to remember to update + the UI as well. And furthermore we have to disable UI editing when the simulation is being updated. But the advantage + is that the UI becomes easier to construct, and there is no danger of discrepancies between user data and the UI. + Second, the UI elements themselves are mostly stateless. Instead we keep track of a minimal set of global states, in + particular mouse and keyboard state, section expand/collapse, contents of the text box being edited if any. + + +Automated enable and disable + While each UI item can be set in enabled or disabled state directly, we also provide automation as follows. Each UI + item can be assigned an integer category. Then a :ref:`mjfItemEnable` callback determines whether each category should + be enabled or disabled, based on some program-specific conditions. For example, sliders that can change the values of + MuJoCo model joints should be disabled when the simulation state is being updated. + + + +.. _uiAPI: + +Main API +~~~~~~~~ + +Click on the links below for detailed API reference of the main UI data structures and functions. + + +**Main data structures:** + +- :ref:`mjUI`: An entire UI. +- :ref:`mjuiState`: Global UI state. +- :ref:`mjuiDef`: One entry in the definition table used for UI construction. + + +**Main functions:** + +- :ref:`mjui_update`: Main UI update function. +- :ref:`mjui_render`: Renders the UI. +- :ref:`mjui_event`: Low-level event handler. +- :ref:`mjui_add`: Helper function used to construct a UI. From 643c1248b895457da9ce9b3d492cd41b763c630f Mon Sep 17 00:00:00 2001 From: Alessio Quaglino Date: Mon, 23 Oct 2023 10:00:33 -0700 Subject: [PATCH 16/74] Skip flex edge Jacobian if not needed. When equality constraints and built-in passive forces are not present, edge Jacobians are unused. This gives a 3x speedup in `mj_kinematics` in `floppy_flex` (from 0.2ms/step to 0.06ms/step). Changed the default options of `floppy_flex` to match those of `floppy` for a fair comparison. PiperOrigin-RevId: 575850622 Change-Id: I4c618536fd17ef930153bbd1f1da7d9e0e44abbe --- model/plugin/elasticity/floppy_flex.xml | 4 ++-- src/engine/engine_core_smooth.c | 10 ++++++++++ 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/model/plugin/elasticity/floppy_flex.xml b/model/plugin/elasticity/floppy_flex.xml index a8645f2f..b6a262c3 100644 --- a/model/plugin/elasticity/floppy_flex.xml +++ b/model/plugin/elasticity/floppy_flex.xml @@ -22,7 +22,7 @@ -