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

Adding type suppressions for pyrefly

Browse Source 
PiperOrigin-RevId: 950703735
Change-Id: I6b6d5826a0dd1f4c0a1cbb2a46aee944cd60d884
This commit is contained in:
Hana Joo
2026-07-20 02:39:22 -07:00
committed by Copybara-Service
parent 00abc5465f
commit 78946ca94e
30 changed files with 193 additions and 193 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/_src/collision_convex.py
+8 -8
View File
@@ -56,27 +56,27 @@ def collider(ncon: int):
fn = collision_fn fn = collision_fn
for i in [0, 1]: for i in [0, 1]:
if key.types[i] == GeomType.BOX: if key.types[i] == GeomType.BOX:
infos[i] = mesh.box(infos[i]) infos[i] = mesh.box(infos[i]) # pyrefly: ignore[unsupported-operation]
in_axes[i] = jax.tree_util.tree_map(lambda x: None, infos[i]).replace( in_axes[i] = jax.tree_util.tree_map(lambda x: None, infos[i]).replace(
pos=0, mat=0, size=0, face=0, vert=0 pos=0, mat=0, size=0, face=0, vert=0
) )
elif key.types[i] == GeomType.MESH: elif key.types[i] == GeomType.MESH:
c, cm = infos[i], m._impl.mesh_convex[key.data_ids[i]] c, cm = infos[i], m._impl.mesh_convex[key.data_ids[i]]
infos[i] = ConvexInfo(**vars(c), **vars(cm)) infos[i] = ConvexInfo(**vars(c), **vars(cm)) # pyrefly: ignore[unsupported-operation]
in_axes[i] = jax.tree_util.tree_map(lambda x: None, infos[i]).replace( in_axes[i] = jax.tree_util.tree_map(lambda x: None, infos[i]).replace(
pos=0, mat=0, size=0 pos=0, mat=0, size=0
) )
elif key.types[i] == GeomType.HFIELD: elif key.types[i] == GeomType.HFIELD:
hfield_info = mesh.hfield(m, key.data_ids[i]) hfield_info = mesh.hfield(m, key.data_ids[i])
infos[i] = hfield_info.replace(pos=infos[i].pos, mat=infos[i].mat) infos[i] = hfield_info.replace(pos=infos[i].pos, mat=infos[i].mat) # pyrefly: ignore[unsupported-operation]
in_axes[i] = hfield_info.replace(pos=0, mat=0, data=None) in_axes[i] = hfield_info.replace(pos=0, mat=0, data=None) # pyrefly: ignore[unsupported-operation]
fn = functools.partial(fn, subgrid_size=key.subgrid_size) fn = functools.partial(fn, subgrid_size=key.subgrid_size)
dist, pos, frame = jax.vmap(fn, in_axes=in_axes)(*infos) # pytype: disable=wrong-keyword-args dist, pos, frame = jax.vmap(fn, in_axes=in_axes)(*infos) # pytype: disable=wrong-keyword-args
if ncon > 1: if ncon > 1:
return jax.tree_util.tree_map(jp.concatenate, (dist, pos, frame)) return jax.tree_util.tree_map(jp.concatenate, (dist, pos, frame))
return dist, pos, frame return dist, pos, frame
collide.ncon = ncon collide.ncon = ncon # pyrefly: ignore[missing-attribute]
return collide return collide
return wrapper return wrapper
@@ -1051,7 +1051,7 @@ def hfield_sphere(
) -> Collision: ) -> Collision:
"""Calculates contacts between a hfield and a sphere.""" """Calculates contacts between a hfield and a sphere."""
rbound = jp.max(s.size) rbound = jp.max(s.size)
dist, pos, n = _hfield_collision(_sphere_convex, h, s, rbound, subgrid_size) dist, pos, n = _hfield_collision(_sphere_convex, h, s, rbound, subgrid_size) # pyrefly: ignore[bad-argument-type]
n_mean = jp.mean(n, axis=0) n_mean = jp.mean(n, axis=0)
mask = dist < jp.minimum(0, dist.min() + 1e-3) mask = dist < jp.minimum(0, dist.min() + 1e-3)
@@ -1075,7 +1075,7 @@ def hfield_capsule(
) -> Collision: ) -> Collision:
"""Calculates contacts between a hfield and a capsule.""" """Calculates contacts between a hfield and a capsule."""
rbound = c.size[0] + c.size[1] rbound = c.size[0] + c.size[1]
dist, pos, n = _hfield_collision(_capsule_convex, h, c, rbound, subgrid_size) dist, pos, n = _hfield_collision(_capsule_convex, h, c, rbound, subgrid_size) # pyrefly: ignore[bad-argument-type]
n_mean = jp.mean(n, axis=0) n_mean = jp.mean(n, axis=0)
mask = dist < jp.minimum(0, dist.min() + 1e-3) mask = dist < jp.minimum(0, dist.min() + 1e-3)
@@ -1099,7 +1099,7 @@ def hfield_convex(
) -> Collision: ) -> Collision:
"""Calculates contacts between a hfield and a capsule.""" """Calculates contacts between a hfield and a capsule."""
rbound = jp.max(c.size) rbound = jp.max(c.size)
dist, pos, n = _hfield_collision(_convex_convex, h, c, rbound, subgrid_size) dist, pos, n = _hfield_collision(_convex_convex, h, c, rbound, subgrid_size) # pyrefly: ignore[bad-argument-type]
n_mean = jp.mean(n, axis=0) n_mean = jp.mean(n, axis=0)
mask = dist < jp.minimum(0, dist.min() + 1e-3) mask = dist < jp.minimum(0, dist.min() + 1e-3)
mjx/mujoco/mjx/_src/collision_driver.py
+3 -3
View File
@@ -384,7 +384,7 @@ def make_condim(
if k.types[1] == mujoco.mjtGeom.mjGEOM_SDF: if k.types[1] == mujoco.mjtGeom.mjGEOM_SDF:
ncon = sdf_initpoints ncon = sdf_initpoints
else: else:
func = _COLLISION_FUNC.get(k.types, None) func = _COLLISION_FUNC.get(k.types, None) # pyrefly: ignore[no-matching-overload]
if func is not None: if func is not None:
ncon = func.ncon # pytype: disable=attribute-error ncon = func.ncon # pytype: disable=attribute-error
else: else:
@@ -430,7 +430,7 @@ def collision(m: Model, d: Data) -> Data:
contact = jax.tree_util.tree_map(lambda x, idx=idx: x[idx], contact) contact = jax.tree_util.tree_map(lambda x, idx=idx: x[idx], contact)
# run the collision function specified by the grouping key # run the collision function specified by the grouping key
func = _COLLISION_FUNC[key.types] func = _COLLISION_FUNC[key.types] # pyrefly: ignore[bad-index]
ncon = func.ncon # pytype: disable=attribute-error ncon = func.ncon # pytype: disable=attribute-error
dist, pos, frame = func(m, d, key, contact.geom) dist, pos, frame = func(m, d, key, contact.geom)
@@ -457,4 +457,4 @@ def collision(m: Model, d: Data) -> Data:
contacts = sum([condim_groups[k] for k in sorted(condim_groups)], []) contacts = sum([condim_groups[k] for k in sorted(condim_groups)], [])
contact = jax.tree_util.tree_map(lambda *x: jp.concatenate(x), *contacts) contact = jax.tree_util.tree_map(lambda *x: jp.concatenate(x), *contacts)
return d.tree_replace({'_impl.contact': contact}) return d.tree_replace({'_impl.contact': contact}) # pyrefly: ignore[bad-return]
mjx/mujoco/mjx/_src/collision_primitive.py
+1 -1
View File
@@ -40,7 +40,7 @@ def collider(ncon: int):
return jax.tree_util.tree_map(jp.concatenate, (dist, pos, frame)) return jax.tree_util.tree_map(jp.concatenate, (dist, pos, frame))
return dist, pos, frame return dist, pos, frame
collide.ncon = ncon collide.ncon = ncon # pyrefly: ignore[missing-attribute]
return collide return collide
return wrapper return wrapper
mjx/mujoco/mjx/_src/collision_sdf.py
+1 -1
View File
@@ -54,7 +54,7 @@ def collider(ncon: int):
return jax.tree_util.tree_map(jp.concatenate, (dist, pos, frame)) return jax.tree_util.tree_map(jp.concatenate, (dist, pos, frame))
return dist, pos, frame return dist, pos, frame
collide.ncon = ncon collide.ncon = ncon # pyrefly: ignore[missing-attribute]
return collide return collide
return wrapper return wrapper
mjx/mujoco/mjx/_src/constraint.py
+4 -4
View File
@@ -716,8 +716,8 @@ def make_constraint(m: Model, d: Data) -> Data:
if not efcs: if not efcs:
z = jp.empty(0) z = jp.empty(0)
d = d.tree_replace({'_impl.efc_J': jp.empty((0, m.nv))}) d = d.tree_replace({'_impl.efc_J': jp.empty((0, m.nv))}) # pyrefly: ignore[bad-assignment]
d = d.tree_replace({ d = d.tree_replace({ # pyrefly: ignore[bad-assignment]
'_impl.efc_D': z, '_impl.efc_D': z,
'_impl.efc_aref': z, '_impl.efc_aref': z,
'_impl.efc_frictionloss': z, '_impl.efc_frictionloss': z,
@@ -736,13 +736,13 @@ def make_constraint(m: Model, d: Data) -> Data:
return aref, r, efc.pos_aref + efc.margin, efc.margin, efc.frictionloss return aref, r, efc.pos_aref + efc.margin, efc.margin, efc.frictionloss
aref, r, pos, margin, frictionloss = fn(efc) aref, r, pos, margin, frictionloss = fn(efc)
d = d.tree_replace({ d = d.tree_replace({ # pyrefly: ignore[bad-assignment]
'_impl.efc_J': efc.J, '_impl.efc_J': efc.J,
'_impl.efc_D': 1 / r, '_impl.efc_D': 1 / r,
'_impl.efc_aref': aref, '_impl.efc_aref': aref,
'_impl.efc_pos': pos, '_impl.efc_pos': pos,
'_impl.efc_margin': margin, '_impl.efc_margin': margin,
}) })
d = d.tree_replace({'_impl.efc_frictionloss': frictionloss}) d = d.tree_replace({'_impl.efc_frictionloss': frictionloss}) # pyrefly: ignore[bad-assignment]
return d return d
mjx/mujoco/mjx/_src/dataclasses.py
+4 -4
View File
@@ -75,7 +75,7 @@ def dataclass(clz: _T, register_as_pytree: bool) -> _T:
Returns: Returns:
the resulting dataclass, registered with Jax the resulting dataclass, registered with Jax
""" """
data_clz = dataclasses.dataclass(frozen=True)(clz) data_clz = dataclasses.dataclass(frozen=True)(clz) # pyrefly: ignore[bad-argument-type]
data_clz.replace = dataclasses.replace data_clz.replace = dataclasses.replace
if register_as_pytree: if register_as_pytree:
@@ -137,7 +137,7 @@ def dataclass(clz: _T, register_as_pytree: bool) -> _T:
data_clz, iterate_clz_with_keys, clz_from_iterable data_clz, iterate_clz_with_keys, clz_from_iterable
) )
return data_clz return data_clz # pyrefly: ignore[bad-return]
TNode = TypeVar('TNode', bound='PyTreeNode') TNode = TypeVar('TNode', bound='PyTreeNode')
@@ -163,7 +163,7 @@ class PyTreeNode:
@classmethod @classmethod
def fields(cls) -> Tuple[dataclasses.Field, ...]: # pylint: disable=g-bare-generic def fields(cls) -> Tuple[dataclasses.Field, ...]: # pylint: disable=g-bare-generic
return dataclasses.fields(cls) return dataclasses.fields(cls) # pyrefly: ignore[bad-argument-type]
def tree_replace( def tree_replace(
self, params: Dict[str, Optional[jax.typing.ArrayLike]] self, params: Dict[str, Optional[jax.typing.ArrayLike]]
@@ -190,7 +190,7 @@ def _tree_replace(
for i, g in enumerate(lst): for i, g in enumerate(lst):
if not hasattr(g, attr[1]): if not hasattr(g, attr[1]):
continue continue
v = val if not hasattr(val, '__iter__') else val[i] v = val if not hasattr(val, '__iter__') else val[i] # pyrefly: ignore[bad-index, unsupported-operation]
lst[i] = _tree_replace(g, attr[1:], v) lst[i] = _tree_replace(g, attr[1:], v)
return base.replace(**{attr[0]: lst}) return base.replace(**{attr[0]: lst})
mjx/mujoco/mjx/_src/forward.py
+2 -2
View File
@@ -89,7 +89,7 @@ def fwd_velocity(m: Model, d: Data) -> Data:
if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX): if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
raise ValueError('fwd_velocity requires JAX backend implementation.') raise ValueError('fwd_velocity requires JAX backend implementation.')
d = d.tree_replace({ d = d.tree_replace({ # pyrefly: ignore[bad-assignment]
'_impl.actuator_velocity': d._impl.actuator_moment @ d.qvel, '_impl.actuator_velocity': d._impl.actuator_moment @ d.qvel,
'_impl.ten_velocity': d._impl.ten_J @ d.qvel, '_impl.ten_velocity': d._impl.ten_J @ d.qvel,
}) })
@@ -357,7 +357,7 @@ def euler(m: Model, d: Data) -> Data:
else: else:
M = d._impl.M + jp.diag(m.opt.timestep * m.dof_damping) M = d._impl.M + jp.diag(m.opt.timestep * m.dof_damping)
dh = d.tree_replace({'_impl.M': M}) dh = d.tree_replace({'_impl.M': M})
dh = smooth.factor_m(m, dh) dh = smooth.factor_m(m, dh) # pyrefly: ignore[bad-argument-type]
qfrc = d.qfrc_smooth + d.qfrc_constraint qfrc = d.qfrc_smooth + d.qfrc_constraint
qacc = smooth.solve_m(m, dh, qfrc) qacc = smooth.solve_m(m, dh, qfrc)
return _advance(m, d, d.act_dot, qacc) return _advance(m, d, d.act_dot, qacc)
mjx/mujoco/mjx/_src/inverse.py
+1 -1
View File
@@ -76,7 +76,7 @@ def inv_constraint(m: Model, d: Data) -> Data:
# update # update
ctx = solver.Context.create(m, d, grad=False) ctx = solver.Context.create(m, d, grad=False)
return d.tree_replace({ return d.tree_replace({ # pyrefly: ignore[bad-return]
'qfrc_constraint': ctx.qfrc_constraint, 'qfrc_constraint': ctx.qfrc_constraint,
'_impl.efc_force': ctx.efc_force, '_impl.efc_force': ctx.efc_force,
}) })
mjx/mujoco/mjx/_src/io.py
+25 -25
View File
@@ -151,15 +151,15 @@ def _resolve_impl_and_device(
if (has_impl, has_device) == (True, True): if (has_impl, has_device) == (True, True):
pass pass
elif (has_impl, has_device) == (True, False): elif (has_impl, has_device) == (True, False):
device = _resolve_device(impl) device = _resolve_device(impl) # pyrefly: ignore[bad-argument-type]
elif (has_impl, has_device) == (False, True): elif (has_impl, has_device) == (False, True):
impl = _resolve_impl(device) impl = _resolve_impl(device) # pyrefly: ignore[bad-argument-type]
else: else:
device = jax.devices()[0] device = jax.devices()[0]
logging.info('Using JAX default device: %s.', device) logging.info('Using JAX default device: %s.', device)
impl = _resolve_impl(device) impl = _resolve_impl(device)
_check_impl_device_compatibility(impl, device) _check_impl_device_compatibility(impl, device) # pyrefly: ignore[bad-argument-type]
return impl, device # pytype: disable=bad-return-type return impl, device # pytype: disable=bad-return-type
@@ -724,7 +724,7 @@ def _make_data_jax(
qpos=jp.array(m.qpos0, dtype=float_), qpos=jp.array(m.qpos0, dtype=float_),
eq_active=m.eq_active0, eq_active=m.eq_active0,
_impl=impl, _impl=impl,
**_make_data_public_fields(m), **_make_data_public_fields(m), # pyrefly: ignore[bad-argument-type]
) )
if m.nmocap: if m.nmocap:
@@ -773,7 +773,7 @@ def _make_data_warp(
nvmax=nvmax, nvmax=nvmax,
) # pylint: disable=undefined-variable ) # pylint: disable=undefined-variable
fields = _make_data_public_fields(m) fields = _make_data_public_fields(m) # pyrefly: ignore[bad-argument-type]
for k in fields: for k in fields:
if k in {'userdata', 'plugin_state', 'history'}: if k in {'userdata', 'plugin_state', 'history'}:
continue continue
@@ -1287,14 +1287,14 @@ def _get_data_into_warp(
if batched if batched
else d else d
) )
result_i = result[i] if batched else result result_i = result[i] if batched else result # pyrefly: ignore[bad-index]
ncon = d_i._impl.nacon[0] ncon = d_i._impl.nacon[0]
nefc = int(d_i._impl.nefc) nefc = int(d_i._impl.nefc)
# nj = int(d_i._impl.nj[0]) # nj = int(d_i._impl.nj[0])
nj = 0 # TODO(btaba): add nj back nj = 0 # TODO(btaba): add nj back
if ncon != result_i.ncon or nefc != result_i.nefc or nj != result_i.nJ: if ncon != result_i.ncon or nefc != result_i.nefc or nj != result_i.nJ: # pyrefly: ignore[missing-attribute]
mujoco._functions._realloc_con_efc(result_i, ncon=ncon, nefc=nefc, nJ=nj) # pylint: disable=protected-access mujoco._functions._realloc_con_efc(result_i, ncon=ncon, nefc=nefc, nJ=nj) # pylint: disable=protected-access # pyrefly: ignore[bad-argument-type]
all_fields = types.Data.fields() + mjxw.types.DataWarp.fields() all_fields = types.Data.fields() + mjxw.types.DataWarp.fields()
for field in all_fields: for field in all_fields:
@@ -1368,14 +1368,14 @@ def _get_data_into(
for i in range(batch_size): for i in range(batch_size):
d_i = jax.tree_util.tree_map(lambda x, i=i: x[i], d) if batched else d d_i = jax.tree_util.tree_map(lambda x, i=i: x[i], d) if batched else d
result_i = result[i] if batched else result result_i = result[i] if batched else result # pyrefly: ignore[bad-index]
ncon = (d_i._impl.contact.dist <= 0).sum() ncon = (d_i._impl.contact.dist <= 0).sum()
efc_active = (d_i._impl.efc_J != 0).any(axis=1) efc_active = (d_i._impl.efc_J != 0).any(axis=1)
nefc = int(efc_active.sum()) nefc = int(efc_active.sum())
nj = (d_i._impl.efc_J != 0).sum() if support.is_sparse(m) else nefc * m.nv nj = (d_i._impl.efc_J != 0).sum() if support.is_sparse(m) else nefc * m.nv
if ncon != result_i.ncon or nefc != result_i.nefc or nj != result_i.nJ: if ncon != result_i.ncon or nefc != result_i.nefc or nj != result_i.nJ: # pyrefly: ignore[missing-attribute]
mujoco._functions._realloc_con_efc(result_i, ncon=ncon, nefc=nefc, nJ=nj) # pylint: disable=protected-access mujoco._functions._realloc_con_efc(result_i, ncon=ncon, nefc=nefc, nJ=nj) # pylint: disable=protected-access # pyrefly: ignore[bad-argument-type]
if d.impl == types.Impl.JAX: if d.impl == types.Impl.JAX:
all_fields = types.Data.fields() + types.DataJAX.fields() all_fields = types.Data.fields() + types.DataJAX.fields()
@@ -1389,10 +1389,10 @@ def _get_data_into(
continue continue
if field.name == 'contact': if field.name == 'contact':
_get_contact(result_i.contact, d_i._impl.contact) _get_contact(result_i.contact, d_i._impl.contact) # pyrefly: ignore[missing-attribute]
# efc_address must be updated because rows were deleted above: # efc_address must be updated because rows were deleted above:
efc_map = np.cumsum(efc_active) - 1 efc_map = np.cumsum(efc_active) - 1
result_i.contact.efc_address[:] = efc_map[result_i.contact.efc_address] result_i.contact.efc_address[:] = efc_map[result_i.contact.efc_address] # pyrefly: ignore[missing-attribute]
continue continue
# MuJoCo actuator_moment is sparse, MJX uses a dense representation. # MuJoCo actuator_moment is sparse, MJX uses a dense representation.
@@ -1412,10 +1412,10 @@ def _get_data_into(
) )
else: else:
actuator_moment = d_i._impl.actuator_moment actuator_moment = d_i._impl.actuator_moment
result_i.moment_rownnz[:] = moment_rownnz result_i.moment_rownnz[:] = moment_rownnz # pyrefly: ignore[missing-attribute]
result_i.moment_rowadr[:] = moment_rowadr result_i.moment_rowadr[:] = moment_rowadr # pyrefly: ignore[missing-attribute]
result_i.moment_colind[:] = moment_colind result_i.moment_colind[:] = moment_colind # pyrefly: ignore[missing-attribute]
result_i.actuator_moment[:] = actuator_moment result_i.actuator_moment[:] = actuator_moment # pyrefly: ignore[missing-attribute]
continue continue
# MuJoCo ten_J is sparse, MJX uses a dense representation. # MuJoCo ten_J is sparse, MJX uses a dense representation.
@@ -1435,7 +1435,7 @@ def _get_data_into(
) )
else: else:
ten_j = d_i._impl.ten_J ten_j = d_i._impl.ten_J
result_i.ten_J[:] = ten_j result_i.ten_J[:] = ten_j # pyrefly: ignore[missing-attribute]
continue continue
if hasattr(d_i._impl, field.name): if hasattr(d_i._impl, field.name):
@@ -1461,9 +1461,9 @@ def _get_data_into(
efc_J_rowadr, efc_J_rowadr,
efc_J_colind, efc_J_colind,
) )
result_i.efc_J_rownnz[:] = efc_J_rownnz result_i.efc_J_rownnz[:] = efc_J_rownnz # pyrefly: ignore[missing-attribute]
result_i.efc_J_rowadr[:] = efc_J_rowadr result_i.efc_J_rowadr[:] = efc_J_rowadr # pyrefly: ignore[missing-attribute]
result_i.efc_J_colind[:] = efc_J_colind result_i.efc_J_colind[:] = efc_J_colind # pyrefly: ignore[missing-attribute]
value = efc_J value = efc_J
else: else:
value = value.reshape(-1) value = value.reshape(-1)
@@ -1496,7 +1496,7 @@ def _get_data_into(
# recalculate qLD and qLDiagInv as MJX and MuJoCo have different # recalculate qLD and qLDiagInv as MJX and MuJoCo have different
# representations of the Cholesky decomposition. # representations of the Cholesky decomposition.
mujoco.mj_factorM(m, result_i) mujoco.mj_factorM(m, result_i) # pyrefly: ignore[bad-argument-type]
# TODO(josechenf): Iterate on the keepalive implementation to make it easier to # TODO(josechenf): Iterate on the keepalive implementation to make it easier to
@@ -1533,7 +1533,7 @@ def _get_data_into_cpp(
d_i: types.Data = ( d_i: types.Data = (
jax.tree_util.tree_map(lambda x, i=i: x[i], d) if batched else d jax.tree_util.tree_map(lambda x, i=i: x[i], d) if batched else d
) )
result_i = result[i] if batched else result result_i = result[i] if batched else result # pyrefly: ignore[bad-index]
if batched: if batched:
addr_i = int(d_impl.pointer_lo[i]) | (int(d_impl.pointer_hi[i]) << 32) addr_i = int(d_impl.pointer_lo[i]) | (int(d_impl.pointer_hi[i]) << 32)
@@ -1562,7 +1562,7 @@ def _get_data_into_cpp(
src_data.mocap_quat[:] = d_i.mocap_quat src_data.mocap_quat[:] = d_i.mocap_quat
mujoco.mj_kinematics(m, src_data) mujoco.mj_kinematics(m, src_data)
mujoco.mj_copyData(result_i, m, src_data) mujoco.mj_copyData(result_i, m, src_data) # pyrefly: ignore[bad-argument-type]
def get_data_into( def get_data_into(
@@ -1810,5 +1810,5 @@ def create_render_context(
_check_warp_installed() _check_warp_installed()
from mujoco.mjx.warp import io as mjxw_io # pylint: disable=g-import-not-at-top # pytype: disable=import-error from mujoco.mjx.warp import io as mjxw_io # pylint: disable=g-import-not-at-top # pytype: disable=import-error
return mjxw_io.create_render_context( return mjxw_io.create_render_context(
mjm, nworld=nworld, devices=devices, **kwargs mjm, nworld=nworld, devices=devices, **kwargs # pyrefly: ignore[bad-argument-type]
) )
mjx/mujoco/mjx/_src/ray.py
+4 -4
View File
@@ -263,8 +263,8 @@ def ray(
for bodyid in bodyexclude: for bodyid in bodyexclude:
geom_filter &= (m.geom_bodyid != bodyid) geom_filter &= (m.geom_bodyid != bodyid)
if geomgroup: if geomgroup:
geomgroup = np.array(geomgroup, dtype=bool) geomgroup = np.array(geomgroup, dtype=bool) # pyrefly: ignore[bad-assignment]
geom_filter &= geomgroup[np.clip(m.geom_group, 0, mujoco.mjNGROUP)] geom_filter &= geomgroup[np.clip(m.geom_group, 0, mujoco.mjNGROUP)] # pyrefly: ignore[bad-index]
# map ray to local geom frames # map ray to local geom frames
geom_pnts = jax.vmap(lambda x, y: x.T @ (pnt - y))(d.geom_xmat, d.geom_xpos) geom_pnts = jax.vmap(lambda x, y: x.T @ (pnt - y))(d.geom_xmat, d.geom_xpos)
@@ -281,7 +281,7 @@ def ray(
args = m.geom_size[id_], geom_pnts[id_], geom_vecs[id_] args = m.geom_size[id_], geom_pnts[id_], geom_vecs[id_]
if geom_type == GeomType.MESH: if geom_type == GeomType.MESH:
dist, id_ = fn(m, id_, *args) dist, id_ = fn(m, id_, *args) # pyrefly: ignore[bad-argument-count, bad-argument-type]
else: else:
dist = jax.vmap(fn)(*args) dist = jax.vmap(fn)(*args)
@@ -314,4 +314,4 @@ def ray_geom(
Returns: Returns:
dist: distance from ray origin to geom surface dist: distance from ray origin to geom surface
""" """
return _RAY_FUNC[geomtype](size, pnt, vec) return _RAY_FUNC[geomtype](size, pnt, vec) # pyrefly: ignore[bad-argument-type, bad-return, missing-argument]
mjx/mujoco/mjx/_src/scan.py
+3 -3
View File
@@ -331,7 +331,7 @@ def flat(
reordered_ys.append(_take(y, _index(ids, input_order))) reordered_ys.append(_take(y, _index(ids, input_order)))
y = reordered_ys if f_ret_is_seq else reordered_ys[0] y = reordered_ys if f_ret_is_seq else reordered_ys[0]
return y return y # pyrefly: ignore[bad-return]
def body_tree( def body_tree(
@@ -443,7 +443,7 @@ def body_tree(
if i < len(in_types): if i < len(in_types):
key_in_take.setdefault(key, []).append(ids) key_in_take.setdefault(key, []).append(ids)
else: else:
key_y_take.setdefault(key, []).append(np.hstack(ids)) key_y_take.setdefault(key, []).append(np.hstack(ids)) # pyrefly: ignore[no-matching-overload]
# use this grouping to take the right data subsets and call vmap(f) # use this grouping to take the right data subsets and call vmap(f)
keys = sorted(key_body_ids, reverse=reverse) keys = sorted(key_body_ids, reverse=reverse)
@@ -494,4 +494,4 @@ def body_tree(
y = y[0] if len(out_types) == 1 else y y = y[0] if len(out_types) == 1 else y
return y return y # pyrefly: ignore[bad-return]
mjx/mujoco/mjx/_src/sensor.py
+12 -12
View File
@@ -339,9 +339,9 @@ def sensor_vel(m: Model, d: Data) -> Data:
elif sensor_type == SensorType.JOINTVEL: elif sensor_type == SensorType.JOINTVEL:
sensor = d.qvel[m.jnt_dofadr[objid]] sensor = d.qvel[m.jnt_dofadr[objid]]
elif sensor_type == SensorType.TENDONVEL: elif sensor_type == SensorType.TENDONVEL:
sensor = d._impl.ten_velocity[objid] sensor = d._impl.ten_velocity[objid] # pyrefly: ignore[missing-attribute]
elif sensor_type == SensorType.ACTUATORVEL: elif sensor_type == SensorType.ACTUATORVEL:
sensor = d._impl.actuator_velocity[objid] sensor = d._impl.actuator_velocity[objid] # pyrefly: ignore[missing-attribute]
elif sensor_type == SensorType.BALLANGVEL: elif sensor_type == SensorType.BALLANGVEL:
jnt_dotadr = m.jnt_dofadr[objid, None] + np.arange(3)[None] jnt_dotadr = m.jnt_dofadr[objid, None] + np.arange(3)[None]
sensor = d.qvel[jnt_dotadr] sensor = d.qvel[jnt_dotadr]
@@ -403,10 +403,10 @@ def sensor_vel(m: Model, d: Data) -> Data:
adrs.append(adrt.reshape(-1)) adrs.append(adrt.reshape(-1))
continue # avoid adding to sensors/adrs list a second time continue # avoid adding to sensors/adrs list a second time
elif sensor_type == SensorType.SUBTREELINVEL: elif sensor_type == SensorType.SUBTREELINVEL:
sensor = d._impl.subtree_linvel[objid] sensor = d._impl.subtree_linvel[objid] # pyrefly: ignore[missing-attribute]
adr = (adr[:, None] + np.arange(3)[None]).reshape(-1) adr = (adr[:, None] + np.arange(3)[None]).reshape(-1)
elif sensor_type == SensorType.SUBTREEANGMOM: elif sensor_type == SensorType.SUBTREEANGMOM:
sensor = d._impl.subtree_angmom[objid] sensor = d._impl.subtree_angmom[objid] # pyrefly: ignore[missing-attribute]
adr = (adr[:, None] + np.arange(3)[None]).reshape(-1) adr = (adr[:, None] + np.arange(3)[None]).reshape(-1)
else: else:
# TODO(taylorhowell): raise error after adding sensor check to io.py # TODO(taylorhowell): raise error after adding sensor check to io.py
@@ -502,7 +502,7 @@ def sensor_acc(m: Model, d: Data) -> Data:
# compute conray, flip if second body # compute conray, flip if second body
conray = jax.vmap( conray = jax.vmap(
lambda frame, force: math.normalize(frame[0] * force[0]) lambda frame, force: math.normalize(frame[0] * force[0])
)(d._impl.contact.frame, contact_force) )(d._impl.contact.frame, contact_force) # pyrefly: ignore[unbound-name]
conray = jp.where(conbody1[..., None], -conray, conray) conray = jp.where(conbody1[..., None], -conray, conray)
# compute distance, mapping over sites and contacts # compute distance, mapping over sites and contacts
@@ -547,7 +547,7 @@ def sensor_acc(m: Model, d: Data) -> Data:
# compute force magnitude for each contact # compute force magnitude for each contact
force_mag = jax.vmap( force_mag = jax.vmap(
lambda forcetorque: jp.dot(forcetorque[:3], forcetorque[:3]) lambda forcetorque: jp.dot(forcetorque[:3], forcetorque[:3])
)(contact_force) )(contact_force) # pyrefly: ignore[unbound-name]
def _reduce(reduction, mask): def _reduce(reduction, mask):
if reduction == 1: # mindist if reduction == 1: # mindist
@@ -662,7 +662,7 @@ def sensor_acc(m: Model, d: Data) -> Data:
slot.append(jp.repeat(nfound, num)[:, None]) slot.append(jp.repeat(nfound, num)[:, None])
if dataspec & (1 << 1): # force if dataspec & (1 << 1): # force
slot.append(flip * contact_force[cid, :3]) slot.append(flip * contact_force[cid, :3]) # pyrefly: ignore[unbound-name]
if dataspec & (1 << 2): # torque if dataspec & (1 << 2): # torque
slot.append(flip * contact_force[cid, 3:]) slot.append(flip * contact_force[cid, 3:])
@@ -699,14 +699,14 @@ def sensor_acc(m: Model, d: Data) -> Data:
bodyid = m.site_bodyid[objid] bodyid = m.site_bodyid[objid]
rot = d.site_xmat[objid] rot = d.site_xmat[objid]
cvel = d.cvel[bodyid] cvel = d.cvel[bodyid]
cacc = d._impl.cacc[bodyid] cacc = d._impl.cacc[bodyid] # pyrefly: ignore[missing-attribute]
dif = d.site_xpos[objid] - d.subtree_com[m.body_rootid[bodyid]] dif = d.site_xpos[objid] - d.subtree_com[m.body_rootid[bodyid]]
sensor = _accelerometer(cvel, cacc, dif, rot) sensor = _accelerometer(cvel, cacc, dif, rot)
adr = (adr[:, None] + np.arange(3)[None]).reshape(-1) adr = (adr[:, None] + np.arange(3)[None]).reshape(-1)
elif sensor_type == SensorType.FORCE: elif sensor_type == SensorType.FORCE:
bodyid = m.site_bodyid[objid] bodyid = m.site_bodyid[objid]
cfrc_int = d._impl.cfrc_int[bodyid] cfrc_int = d._impl.cfrc_int[bodyid] # pyrefly: ignore[missing-attribute]
site_xmat = d.site_xmat[objid] site_xmat = d.site_xmat[objid]
sensor = jax.vmap(lambda mat, vec: mat.T @ vec)( sensor = jax.vmap(lambda mat, vec: mat.T @ vec)(
site_xmat, cfrc_int[:, 3:] site_xmat, cfrc_int[:, 3:]
@@ -715,7 +715,7 @@ def sensor_acc(m: Model, d: Data) -> Data:
elif sensor_type == SensorType.TORQUE: elif sensor_type == SensorType.TORQUE:
bodyid = m.site_bodyid[objid] bodyid = m.site_bodyid[objid]
rootid = m.body_rootid[bodyid] rootid = m.body_rootid[bodyid]
cfrc_int = d._impl.cfrc_int[bodyid] cfrc_int = d._impl.cfrc_int[bodyid] # pyrefly: ignore[missing-attribute]
site_xmat = d.site_xmat[objid] site_xmat = d.site_xmat[objid]
dif = d.site_xpos[objid] - d.subtree_com[rootid] dif = d.site_xpos[objid] - d.subtree_com[rootid]
sensor = jax.vmap( sensor = jax.vmap(
@@ -744,7 +744,7 @@ def sensor_acc(m: Model, d: Data) -> Data:
pos, bodyid = objtype_data[ot] pos, bodyid = objtype_data[ot]
pos = pos[objidt] pos = pos[objidt]
bodyid = bodyid[objidt] bodyid = bodyid[objidt]
cacc = d._impl.cacc[bodyid] cacc = d._impl.cacc[bodyid] # pyrefly: ignore[missing-attribute]
if sensor_type == SensorType.FRAMELINACC: if sensor_type == SensorType.FRAMELINACC:
@@ -774,7 +774,7 @@ def sensor_acc(m: Model, d: Data) -> Data:
# TODO(taylorhowell): raise error after adding sensor check to io.py # TODO(taylorhowell): raise error after adding sensor check to io.py
continue # unsupported sensor type continue # unsupported sensor type
sensors.append(_apply_cutoff(sensor, cutoff, data_type[0]).reshape(-1)) sensors.append(_apply_cutoff(sensor, cutoff, data_type[0]).reshape(-1)) # pyrefly: ignore[bad-argument-type]
adrs.append(adr) adrs.append(adr)
if not adrs: if not adrs:
mjx/mujoco/mjx/_src/smooth.py
+14 -14
View File
@@ -173,7 +173,7 @@ def com_pos(m: Model, d: Data) -> Data:
root_com = subtree_com[m.body_rootid] root_com = subtree_com[m.body_rootid]
offset = d.xipos - root_com offset = d.xipos - root_com
cinert = inert_com(m.body_inertia, d.ximat, offset, m.body_mass) cinert = inert_com(m.body_inertia, d.ximat, offset, m.body_mass)
d = d.tree_replace({'_impl.cinert': cinert}) d = d.tree_replace({'_impl.cinert': cinert}) # pyrefly: ignore[bad-assignment]
# map motion dofs to global frame centered at subtree_com # map motion dofs to global frame centered at subtree_com
def cdof_fn(jnt_typs, root_com, xmat, xanchor, xaxis): def cdof_fn(jnt_typs, root_com, xmat, xanchor, xaxis):
@@ -212,7 +212,7 @@ def com_pos(m: Model, d: Data) -> Data:
d.xanchor, d.xanchor,
d.xaxis, d.xaxis,
) )
d = d.tree_replace({'cdof': cdof}) d = d.tree_replace({'cdof': cdof}) # pyrefly: ignore[bad-assignment]
return d return d
@@ -302,12 +302,12 @@ def crb(m: Model, d: Data) -> Data:
crb_body = scan.body_tree(m, crb_fn, 'b', 'b', d._impl.cinert, reverse=True) crb_body = scan.body_tree(m, crb_fn, 'b', 'b', d._impl.cinert, reverse=True)
crb_body = crb_body.at[0].set(0.0) crb_body = crb_body.at[0].set(0.0)
d = d.tree_replace({'_impl.crb': crb_body}) d = d.tree_replace({'_impl.crb': crb_body}) # pyrefly: ignore[bad-assignment]
crb_dof = jp.take(crb_body, jp.array(m.dof_bodyid), axis=0) crb_dof = jp.take(crb_body, jp.array(m.dof_bodyid), axis=0)
crb_cdof = jax.vmap(math.inert_mul)(crb_dof, d.cdof) crb_cdof = jax.vmap(math.inert_mul)(crb_dof, d.cdof)
qm = support.make_m(m, crb_cdof, d.cdof, m.dof_armature) qm = support.make_m(m, crb_cdof, d.cdof, m.dof_armature)
d = d.tree_replace({'_impl.M': qm}) d = d.tree_replace({'_impl.M': qm}) # pyrefly: ignore[bad-assignment]
return d return d
@@ -318,7 +318,7 @@ def factor_m(m: Model, d: Data) -> Data:
if not support.is_sparse(m): if not support.is_sparse(m):
qh, _ = jax.scipy.linalg.cho_factor(d._impl.M) qh, _ = jax.scipy.linalg.cho_factor(d._impl.M)
d = d.tree_replace({'_impl.qLD': qh}) d = d.tree_replace({'_impl.qLD': qh}) # pyrefly: ignore[bad-assignment]
return d return d
depth = [] depth = []
@@ -367,7 +367,7 @@ def factor_m(m: Model, d: Data) -> Data:
qld_diag = qld[diag_adr] qld_diag = qld[diag_adr]
qld = (qld / qld[np.array(diag_ds)]).at[diag_adr].set(qld_diag) qld = (qld / qld[np.array(diag_ds)]).at[diag_adr].set(qld_diag)
d = d.tree_replace({'_impl.qLD': qld, '_impl.qLDiagInv': 1 / qld_diag}) d = d.tree_replace({'_impl.qLD': qld, '_impl.qLDiagInv': 1 / qld_diag}) # pyrefly: ignore[bad-assignment]
return d return d
@@ -448,7 +448,7 @@ def com_vel(m: Model, d: Data) -> Data:
d.qvel, d.qvel,
) )
d = d.tree_replace({'cvel': cvel, 'cdof_dot': cdof_dot}) d = d.tree_replace({'cvel': cvel, 'cdof_dot': cdof_dot}) # pyrefly: ignore[bad-assignment]
return d return d
@@ -544,7 +544,7 @@ def subtree_vel(m: Model, d: Data) -> Data:
reverse=True, reverse=True,
) )
return d.tree_replace({ return d.tree_replace({ # pyrefly: ignore[bad-return]
'_impl.subtree_linvel': subtree_linvel, '_impl.subtree_linvel': subtree_linvel,
'_impl.subtree_angmom': subtree_angmom, '_impl.subtree_angmom': subtree_angmom,
}) })
@@ -841,7 +841,7 @@ def rne_postconstraint(m: Model, d: Data) -> Data:
) )
# update data # update data
return d.tree_replace({ return d.tree_replace({ # pyrefly: ignore[bad-return]
'_impl.cacc': cacc, '_impl.cacc': cacc,
'_impl.cfrc_int': cfrc_int, '_impl.cfrc_int': cfrc_int,
'_impl.cfrc_ext': cfrc_ext, '_impl.cfrc_ext': cfrc_ext,
@@ -1172,7 +1172,7 @@ def tendon(m: Model, d: Data) -> Data:
[wrap_obj[sort], jp.zeros(2 * m.nwrap - count, dtype=int)] [wrap_obj[sort], jp.zeros(2 * m.nwrap - count, dtype=int)]
).reshape((m.nwrap, 2)) ).reshape((m.nwrap, 2))
return d.tree_replace({ return d.tree_replace({ # pyrefly: ignore[bad-return]
'ten_length': ten_length, 'ten_length': ten_length,
'_impl.ten_J': ten_moment, '_impl.ten_J': ten_moment,
'_impl.ten_wrapadr': jp.array(ten_wrapadr, dtype=int), '_impl.ten_wrapadr': jp.array(ten_wrapadr, dtype=int),
@@ -1276,7 +1276,7 @@ def transmission(m: Model, d: Data) -> Data:
moment = jac @ wrench moment = jac @ wrench
elif trntype == TrnType.TENDON: elif trntype == TrnType.TENDON:
length = d.ten_length[trnid[0]] * gear[:1] length = d.ten_length[trnid[0]] * gear[:1]
moment = d._impl.ten_J[trnid[0]] * gear[0] moment = d._impl.ten_J[trnid[0]] * gear[0] # pyrefly: ignore[missing-attribute]
else: else:
raise RuntimeError(f'unrecognized trntype: {TrnType(trntype)}') raise RuntimeError(f'unrecognized trntype: {TrnType(trntype)}')
@@ -1308,7 +1308,7 @@ def transmission(m: Model, d: Data) -> Data:
length = length.reshape((m.nu,)) length = length.reshape((m.nu,))
moment = moment.reshape((m.nu, m.nv)) moment = moment.reshape((m.nu, m.nv))
d = d.tree_replace( d = d.tree_replace( # pyrefly: ignore[bad-assignment]
{'actuator_length': length, '_impl.actuator_moment': moment} {'actuator_length': length, '_impl.actuator_moment': moment}
) )
return d return d
@@ -1332,7 +1332,7 @@ def tendon_armature(m: Model, d: Data) -> Data:
j = m.M_colind j = m.M_colind
JTAJ = JTAJ[(i, j)] JTAJ = JTAJ[(i, j)]
return d.tree_replace({'_impl.M': d._impl.M + JTAJ}) return d.tree_replace({'_impl.M': d._impl.M + JTAJ}) # pyrefly: ignore[bad-return]
def tendon_dot(m: Model, d: Data) -> jax.Array: def tendon_dot(m: Model, d: Data) -> jax.Array:
@@ -1457,7 +1457,7 @@ def tendon_bias(m: Model, d: Data) -> Data:
# add bias term: qfrc += ten_J * armature * ten_Jdot @ qvel # add bias term: qfrc += ten_J * armature * ten_Jdot @ qvel
coef = m.tendon_armature * jp.dot(ten_Jdot, d.qvel) coef = m.tendon_armature * jp.dot(ten_Jdot, d.qvel)
return d.tree_replace({ return d.tree_replace({ # pyrefly: ignore[bad-return]
'qfrc_bias': ( 'qfrc_bias': (
d.qfrc_bias d.qfrc_bias
+ jp.sum(jax.vmap(jp.multiply)(d._impl.ten_J, coef), axis=0) + jp.sum(jax.vmap(jp.multiply)(d._impl.ten_J, coef), axis=0)
mjx/mujoco/mjx/_src/solver.py
+2 -2
View File
@@ -482,7 +482,7 @@ def _linesearch(m: Model, d: Data, ctx: Context) -> Context:
vv = jp.sum(v[:, 1:] * v[:, 1:], axis=1) vv = jp.sum(v[:, 1:] * v[:, 1:], axis=1)
point_fn = lambda a: _LSPoint.create( point_fn = lambda a: _LSPoint.create(
m, d, ctx, a, jv, quad, quad_gauss, uu, v0, uv, vv m, d, ctx, a, jv, quad, quad_gauss, uu, v0, uv, vv # pyrefly: ignore[bad-argument-type]
) )
def cond(ctx: _LSContext) -> jax.Array: def cond(ctx: _LSContext) -> jax.Array:
@@ -601,7 +601,7 @@ def solve(m: Model, d: Data) -> Data:
else: else:
ctx = jax.lax.while_loop(cond, body, ctx) ctx = jax.lax.while_loop(cond, body, ctx)
d = d.tree_replace({ d = d.tree_replace({ # pyrefly: ignore[bad-assignment]
'qfrc_constraint': ctx.qfrc_constraint, 'qfrc_constraint': ctx.qfrc_constraint,
'qacc': ctx.qacc, 'qacc': ctx.qacc,
'_impl.efc_force': ctx.efc_force, '_impl.efc_force': ctx.efc_force,
mjx/mujoco/mjx/_src/support.py
+6 -6
View File
@@ -501,15 +501,15 @@ class BindData(object):
num = sum((typ == jt) * jt.dof_width() for jt in JointType) num = sum((typ == jt) * jt.dof_width() for jt in JointType)
if isinstance(self.id, list): if isinstance(self.id, list):
idx = [] idx = []
for a, n in zip(adr, num): for a, n in zip(adr, num): # pyrefly: ignore[bad-argument-type]
idx.extend(a + j for j in range(n)) idx.extend(a + j for j in range(n))
return self._slice(self.__getname(name), idx) return self._slice(self.__getname(name), idx)
elif num > 1: elif num > 1:
return self._slice(self.__getname(name), slice(adr, adr + num)) return self._slice(self.__getname(name), slice(adr, adr + num))
else: else:
return self._slice(self.__getname(name), adr) return self._slice(self.__getname(name), adr) # pyrefly: ignore[bad-argument-type]
elif name in ('mocap_pos', 'mocap_quat'): elif name in ('mocap_pos', 'mocap_quat'):
return self._slice(self.__getname(name), self.model.body_mocapid[self.id]) return self._slice(self.__getname(name), self.model.body_mocapid[self.id]) # pyrefly: ignore[bad-argument-type]
return self._slice(self.__getname(name), self.id) return self._slice(self.__getname(name), self.id)
def set(self, name: str, value: jax.Array) -> Data: def set(self, name: str, value: jax.Array) -> Data:
@@ -521,7 +521,7 @@ class BindData(object):
try: try:
iter(value) iter(value)
except TypeError: except TypeError:
value = [value] value = [value] # pyrefly: ignore[bad-assignment]
if name in ('qpos', 'qvel', 'qacc', 'mocap_pos', 'mocap_quat'): if name in ('qpos', 'qvel', 'qacc', 'mocap_pos', 'mocap_quat'):
adr = num = 0 adr = num = 0
if name == 'qpos': if name == 'qpos':
@@ -549,7 +549,7 @@ class BindData(object):
num = [dim] num = [dim]
i = 0 i = 0
value = jax.numpy.array(value).flatten() value = jax.numpy.array(value).flatten()
for a, n in zip(adr, num): for a, n in zip(adr, num): # pyrefly: ignore[bad-argument-type]
shape = array.shape shape = array.shape
array = array.flatten().at[a : a + n].set(value[i : i + n]).reshape(shape) array = array.flatten().at[a : a + n].set(value[i : i + n]).reshape(shape)
i += n i += n
@@ -623,7 +623,7 @@ def contact_force_dim(
if m.opt.cone == ConeType.PYRAMIDAL: if m.opt.cone == ConeType.PYRAMIDAL:
efc_address = ( efc_address = (
d._impl.contact.efc_address[idx_dim, None] # pytype: disable=attribute-error d._impl.contact.efc_address[idx_dim, None] # pytype: disable=attribute-error
+ np.arange(np.where(dim == 1, 1, 2 * (dim - 1)))[None] + np.arange(np.where(dim == 1, 1, 2 * (dim - 1)))[None] # pyrefly: ignore[no-matching-overload]
) )
efc_force = d._impl.efc_force[efc_address] # pytype: disable=attribute-error efc_force = d._impl.efc_force[efc_address] # pytype: disable=attribute-error
force = jax.vmap(_decode_pyramid, in_axes=(0, 0, None))( force = jax.vmap(_decode_pyramid, in_axes=(0, 0, None))(
mjx/mujoco/mjx/_src/test_util.py
+4 -4
View File
@@ -63,13 +63,13 @@ def benchmark(
xla_flags += ' --xla_gpu_triton_gemm_any=True' xla_flags += ' --xla_gpu_triton_gemm_any=True'
os.environ['XLA_FLAGS'] = xla_flags os.environ['XLA_FLAGS'] = xla_flags
m.opt.solver = { m.opt.solver = { # pyrefly: ignore[bad-assignment]
'cg': mujoco.mjtSolver.mjSOL_CG, 'cg': mujoco.mjtSolver.mjSOL_CG,
'newton': mujoco.mjtSolver.mjSOL_NEWTON, 'newton': mujoco.mjtSolver.mjSOL_NEWTON,
}[solver.lower()] }[solver.lower()]
m.opt.iterations = iterations m.opt.iterations = iterations
m.opt.ls_iterations = ls_iterations m.opt.ls_iterations = ls_iterations
m = io.put_model(m) m = io.put_model(m) # pyrefly: ignore[bad-assignment]
@jax.pmap @jax.pmap
def init(key): def init(key):
@@ -449,7 +449,7 @@ def create_mjcf(
if p(80): if p(80):
continue continue
geom1, geom2 = np.random.choice(geom_names, replace=False, size=2) geom1, geom2 = np.random.choice(geom_names, replace=False, size=2) # pyrefly: ignore[no-matching-overload]
if geom1 > geom2: if geom1 > geom2:
geom1, geom2 = geom2, geom1 geom1, geom2 = geom2, geom1
@@ -467,7 +467,7 @@ def create_mjcf(
if p(50): if p(50):
continue continue
body1, body2 = np.random.choice(body_names, replace=False, size=2) body1, body2 = np.random.choice(body_names, replace=False, size=2) # pyrefly: ignore[no-matching-overload]
ET.SubElement(contact, 'exclude', {'body1': body1, 'body2': body2}) ET.SubElement(contact, 'exclude', {'body1': body1, 'body2': body2})
# ElementTree.indent is not available before Python 3.9 # ElementTree.indent is not available before Python 3.9
mjx/mujoco/mjx/_src/types.py
+59 -59
View File
@@ -105,10 +105,10 @@ class JointType(enum.IntEnum):
HINGE: rotation angle (rad) around body-fixed axis (1,) HINGE: rotation angle (rad) around body-fixed axis (1,)
""" """
FREE = mujoco.mjtJoint.mjJNT_FREE FREE = mujoco.mjtJoint.mjJNT_FREE # pyrefly: ignore[bad-assignment]
BALL = mujoco.mjtJoint.mjJNT_BALL BALL = mujoco.mjtJoint.mjJNT_BALL # pyrefly: ignore[bad-assignment]
SLIDE = mujoco.mjtJoint.mjJNT_SLIDE SLIDE = mujoco.mjtJoint.mjJNT_SLIDE # pyrefly: ignore[bad-assignment]
HINGE = mujoco.mjtJoint.mjJNT_HINGE HINGE = mujoco.mjtJoint.mjJNT_HINGE # pyrefly: ignore[bad-assignment]
def dof_width(self) -> int: def dof_width(self) -> int:
return {0: 6, 1: 3, 2: 1, 3: 1}[self.value] return {0: 6, 1: 3, 2: 1, 3: 1}[self.value]
@@ -126,9 +126,9 @@ class IntegratorType(enum.IntEnum):
IMPLICITFAST: implicit in velocity, no rne derivative IMPLICITFAST: implicit in velocity, no rne derivative
""" """
EULER = mujoco.mjtIntegrator.mjINT_EULER EULER = mujoco.mjtIntegrator.mjINT_EULER # pyrefly: ignore[bad-assignment]
RK4 = mujoco.mjtIntegrator.mjINT_RK4 RK4 = mujoco.mjtIntegrator.mjINT_RK4 # pyrefly: ignore[bad-assignment]
IMPLICITFAST = mujoco.mjtIntegrator.mjINT_IMPLICITFAST IMPLICITFAST = mujoco.mjtIntegrator.mjINT_IMPLICITFAST # pyrefly: ignore[bad-assignment]
# unsupported: IMPLICIT # unsupported: IMPLICIT
@@ -147,14 +147,14 @@ class GeomType(enum.IntEnum):
SDF: signed distance field SDF: signed distance field
""" """
PLANE = mujoco.mjtGeom.mjGEOM_PLANE PLANE = mujoco.mjtGeom.mjGEOM_PLANE # pyrefly: ignore[bad-assignment]
HFIELD = mujoco.mjtGeom.mjGEOM_HFIELD HFIELD = mujoco.mjtGeom.mjGEOM_HFIELD # pyrefly: ignore[bad-assignment]
SPHERE = mujoco.mjtGeom.mjGEOM_SPHERE SPHERE = mujoco.mjtGeom.mjGEOM_SPHERE # pyrefly: ignore[bad-assignment]
CAPSULE = mujoco.mjtGeom.mjGEOM_CAPSULE CAPSULE = mujoco.mjtGeom.mjGEOM_CAPSULE # pyrefly: ignore[bad-assignment]
ELLIPSOID = mujoco.mjtGeom.mjGEOM_ELLIPSOID ELLIPSOID = mujoco.mjtGeom.mjGEOM_ELLIPSOID # pyrefly: ignore[bad-assignment]
CYLINDER = mujoco.mjtGeom.mjGEOM_CYLINDER CYLINDER = mujoco.mjtGeom.mjGEOM_CYLINDER # pyrefly: ignore[bad-assignment]
BOX = mujoco.mjtGeom.mjGEOM_BOX BOX = mujoco.mjtGeom.mjGEOM_BOX # pyrefly: ignore[bad-assignment]
MESH = mujoco.mjtGeom.mjGEOM_MESH MESH = mujoco.mjtGeom.mjGEOM_MESH # pyrefly: ignore[bad-assignment]
# unsupported: NGEOMTYPES, ARROW*, LINE, SKIN, LABEL, NONE # unsupported: NGEOMTYPES, ARROW*, LINE, SKIN, LABEL, NONE
@@ -184,8 +184,8 @@ class ConeType(enum.IntEnum):
ELLIPTIC: elliptic ELLIPTIC: elliptic
""" """
PYRAMIDAL = mujoco.mjtCone.mjCONE_PYRAMIDAL PYRAMIDAL = mujoco.mjtCone.mjCONE_PYRAMIDAL # pyrefly: ignore[bad-assignment]
ELLIPTIC = mujoco.mjtCone.mjCONE_ELLIPTIC ELLIPTIC = mujoco.mjtCone.mjCONE_ELLIPTIC # pyrefly: ignore[bad-assignment]
class JacobianType(enum.IntEnum): class JacobianType(enum.IntEnum):
@@ -197,9 +197,9 @@ class JacobianType(enum.IntEnum):
AUTO: sparse if nv>60 and device is TPU, dense otherwise AUTO: sparse if nv>60 and device is TPU, dense otherwise
""" """
DENSE = mujoco.mjtJacobian.mjJAC_DENSE DENSE = mujoco.mjtJacobian.mjJAC_DENSE # pyrefly: ignore[bad-assignment]
SPARSE = mujoco.mjtJacobian.mjJAC_SPARSE SPARSE = mujoco.mjtJacobian.mjJAC_SPARSE # pyrefly: ignore[bad-assignment]
AUTO = mujoco.mjtJacobian.mjJAC_AUTO AUTO = mujoco.mjtJacobian.mjJAC_AUTO # pyrefly: ignore[bad-assignment]
class SolverType(enum.IntEnum): class SolverType(enum.IntEnum):
@@ -211,8 +211,8 @@ class SolverType(enum.IntEnum):
""" """
# unsupported: PGS # unsupported: PGS
CG = mujoco.mjtSolver.mjSOL_CG CG = mujoco.mjtSolver.mjSOL_CG # pyrefly: ignore[bad-assignment]
NEWTON = mujoco.mjtSolver.mjSOL_NEWTON NEWTON = mujoco.mjtSolver.mjSOL_NEWTON # pyrefly: ignore[bad-assignment]
class EqType(enum.IntEnum): class EqType(enum.IntEnum):
@@ -225,10 +225,10 @@ class EqType(enum.IntEnum):
TENDON: couple the lengths of two tendons with cubic TENDON: couple the lengths of two tendons with cubic
""" """
CONNECT = mujoco.mjtEq.mjEQ_CONNECT CONNECT = mujoco.mjtEq.mjEQ_CONNECT # pyrefly: ignore[bad-assignment]
WELD = mujoco.mjtEq.mjEQ_WELD WELD = mujoco.mjtEq.mjEQ_WELD # pyrefly: ignore[bad-assignment]
JOINT = mujoco.mjtEq.mjEQ_JOINT JOINT = mujoco.mjtEq.mjEQ_JOINT # pyrefly: ignore[bad-assignment]
TENDON = mujoco.mjtEq.mjEQ_TENDON TENDON = mujoco.mjtEq.mjEQ_TENDON # pyrefly: ignore[bad-assignment]
# unsupported: DISTANCE # unsupported: DISTANCE
@@ -243,11 +243,11 @@ class WrapType(enum.IntEnum):
CYLINDER: wrap around (infinite) cylinder CYLINDER: wrap around (infinite) cylinder
""" """
JOINT = mujoco.mjtWrap.mjWRAP_JOINT JOINT = mujoco.mjtWrap.mjWRAP_JOINT # pyrefly: ignore[bad-assignment]
PULLEY = mujoco.mjtWrap.mjWRAP_PULLEY PULLEY = mujoco.mjtWrap.mjWRAP_PULLEY # pyrefly: ignore[bad-assignment]
SITE = mujoco.mjtWrap.mjWRAP_SITE SITE = mujoco.mjtWrap.mjWRAP_SITE # pyrefly: ignore[bad-assignment]
SPHERE = mujoco.mjtWrap.mjWRAP_SPHERE SPHERE = mujoco.mjtWrap.mjWRAP_SPHERE # pyrefly: ignore[bad-assignment]
CYLINDER = mujoco.mjtWrap.mjWRAP_CYLINDER CYLINDER = mujoco.mjtWrap.mjWRAP_CYLINDER # pyrefly: ignore[bad-assignment]
class TrnType(enum.IntEnum): class TrnType(enum.IntEnum):
@@ -260,10 +260,10 @@ class TrnType(enum.IntEnum):
SITE: force on site SITE: force on site
""" """
JOINT = mujoco.mjtTrn.mjTRN_JOINT JOINT = mujoco.mjtTrn.mjTRN_JOINT # pyrefly: ignore[bad-assignment]
JOINTINPARENT = mujoco.mjtTrn.mjTRN_JOINTINPARENT JOINTINPARENT = mujoco.mjtTrn.mjTRN_JOINTINPARENT # pyrefly: ignore[bad-assignment]
SITE = mujoco.mjtTrn.mjTRN_SITE SITE = mujoco.mjtTrn.mjTRN_SITE # pyrefly: ignore[bad-assignment]
TENDON = mujoco.mjtTrn.mjTRN_TENDON TENDON = mujoco.mjtTrn.mjTRN_TENDON # pyrefly: ignore[bad-assignment]
# unsupported: SLIDERCRANK, BODY # unsupported: SLIDERCRANK, BODY
@@ -278,11 +278,11 @@ class DynType(enum.IntEnum):
MUSCLE: piece-wise linear filter with two time constants MUSCLE: piece-wise linear filter with two time constants
""" """
NONE = mujoco.mjtDyn.mjDYN_NONE NONE = mujoco.mjtDyn.mjDYN_NONE # pyrefly: ignore[bad-assignment]
INTEGRATOR = mujoco.mjtDyn.mjDYN_INTEGRATOR INTEGRATOR = mujoco.mjtDyn.mjDYN_INTEGRATOR # pyrefly: ignore[bad-assignment]
FILTER = mujoco.mjtDyn.mjDYN_FILTER FILTER = mujoco.mjtDyn.mjDYN_FILTER # pyrefly: ignore[bad-assignment]
FILTEREXACT = mujoco.mjtDyn.mjDYN_FILTEREXACT FILTEREXACT = mujoco.mjtDyn.mjDYN_FILTEREXACT # pyrefly: ignore[bad-assignment]
MUSCLE = mujoco.mjtDyn.mjDYN_MUSCLE MUSCLE = mujoco.mjtDyn.mjDYN_MUSCLE # pyrefly: ignore[bad-assignment]
# unsupported: USER # unsupported: USER
@@ -295,9 +295,9 @@ class GainType(enum.IntEnum):
MUSCLE: muscle FLV curve computed by muscle_gain MUSCLE: muscle FLV curve computed by muscle_gain
""" """
FIXED = mujoco.mjtGain.mjGAIN_FIXED FIXED = mujoco.mjtGain.mjGAIN_FIXED # pyrefly: ignore[bad-assignment]
AFFINE = mujoco.mjtGain.mjGAIN_AFFINE AFFINE = mujoco.mjtGain.mjGAIN_AFFINE # pyrefly: ignore[bad-assignment]
MUSCLE = mujoco.mjtGain.mjGAIN_MUSCLE MUSCLE = mujoco.mjtGain.mjGAIN_MUSCLE # pyrefly: ignore[bad-assignment]
# unsupported: USER # unsupported: USER
@@ -310,9 +310,9 @@ class BiasType(enum.IntEnum):
MUSCLE: muscle passive force computed by muscle_bias MUSCLE: muscle passive force computed by muscle_bias
""" """
NONE = mujoco.mjtBias.mjBIAS_NONE NONE = mujoco.mjtBias.mjBIAS_NONE # pyrefly: ignore[bad-assignment]
AFFINE = mujoco.mjtBias.mjBIAS_AFFINE AFFINE = mujoco.mjtBias.mjBIAS_AFFINE # pyrefly: ignore[bad-assignment]
MUSCLE = mujoco.mjtBias.mjBIAS_MUSCLE MUSCLE = mujoco.mjtBias.mjBIAS_MUSCLE # pyrefly: ignore[bad-assignment]
# unsupported: USER # unsupported: USER
@@ -327,14 +327,14 @@ class ConstraintType(enum.IntEnum):
CONTACT_PYRAMIDAL: frictional contact, pyramidal friction cone CONTACT_PYRAMIDAL: frictional contact, pyramidal friction cone
""" """
EQUALITY = mujoco.mjtConstraint.mjCNSTR_EQUALITY EQUALITY = mujoco.mjtConstraint.mjCNSTR_EQUALITY # pyrefly: ignore[bad-assignment]
FRICTION_DOF = mujoco.mjtConstraint.mjCNSTR_FRICTION_DOF FRICTION_DOF = mujoco.mjtConstraint.mjCNSTR_FRICTION_DOF # pyrefly: ignore[bad-assignment]
FRICTION_TENDON = mujoco.mjtConstraint.mjCNSTR_FRICTION_TENDON FRICTION_TENDON = mujoco.mjtConstraint.mjCNSTR_FRICTION_TENDON # pyrefly: ignore[bad-assignment]
LIMIT_JOINT = mujoco.mjtConstraint.mjCNSTR_LIMIT_JOINT LIMIT_JOINT = mujoco.mjtConstraint.mjCNSTR_LIMIT_JOINT # pyrefly: ignore[bad-assignment]
LIMIT_TENDON = mujoco.mjtConstraint.mjCNSTR_LIMIT_TENDON LIMIT_TENDON = mujoco.mjtConstraint.mjCNSTR_LIMIT_TENDON # pyrefly: ignore[bad-assignment]
CONTACT_FRICTIONLESS = mujoco.mjtConstraint.mjCNSTR_CONTACT_FRICTIONLESS CONTACT_FRICTIONLESS = mujoco.mjtConstraint.mjCNSTR_CONTACT_FRICTIONLESS # pyrefly: ignore[bad-assignment]
CONTACT_PYRAMIDAL = mujoco.mjtConstraint.mjCNSTR_CONTACT_PYRAMIDAL CONTACT_PYRAMIDAL = mujoco.mjtConstraint.mjCNSTR_CONTACT_PYRAMIDAL # pyrefly: ignore[bad-assignment]
CONTACT_ELLIPTIC = mujoco.mjtConstraint.mjCNSTR_CONTACT_ELLIPTIC CONTACT_ELLIPTIC = mujoco.mjtConstraint.mjCNSTR_CONTACT_ELLIPTIC # pyrefly: ignore[bad-assignment]
class CamLightType(enum.IntEnum): class CamLightType(enum.IntEnum):
@@ -348,11 +348,11 @@ class CamLightType(enum.IntEnum):
TARGETBODYCOM: pos fixed in body, rot tracks target subtree com TARGETBODYCOM: pos fixed in body, rot tracks target subtree com
""" """
FIXED = mujoco.mjtCamLight.mjCAMLIGHT_FIXED FIXED = mujoco.mjtCamLight.mjCAMLIGHT_FIXED # pyrefly: ignore[bad-assignment]
TRACK = mujoco.mjtCamLight.mjCAMLIGHT_TRACK TRACK = mujoco.mjtCamLight.mjCAMLIGHT_TRACK # pyrefly: ignore[bad-assignment]
TRACKCOM = mujoco.mjtCamLight.mjCAMLIGHT_TRACKCOM TRACKCOM = mujoco.mjtCamLight.mjCAMLIGHT_TRACKCOM # pyrefly: ignore[bad-assignment]
TARGETBODY = mujoco.mjtCamLight.mjCAMLIGHT_TARGETBODY TARGETBODY = mujoco.mjtCamLight.mjCAMLIGHT_TARGETBODY # pyrefly: ignore[bad-assignment]
TARGETBODYCOM = mujoco.mjtCamLight.mjCAMLIGHT_TARGETBODYCOM TARGETBODYCOM = mujoco.mjtCamLight.mjCAMLIGHT_TARGETBODYCOM # pyrefly: ignore[bad-assignment]
class SensorType(enum.IntEnum): class SensorType(enum.IntEnum):
mjx/mujoco/mjx/codegen/generate_warp_shim.py
+2 -2
View File
@@ -264,8 +264,8 @@ def _warp_function(
# create a dummy output if there are no output fields # create a dummy output if there are no output fields
needs_dummy_output = not field_usage.data_out_fields needs_dummy_output = not field_usage.data_out_fields
if needs_dummy_output and fn_name != 'render': if needs_dummy_output and fn_name != 'render':
fn_args_raw.append('# Dummy output') fn_args_raw.append('# Dummy output') # pyrefly: ignore[bad-argument-type]
fn_args_raw.append('dummy: wp.array[int],') fn_args_raw.append('dummy: wp.array[int],') # pyrefly: ignore[bad-argument-type]
return fn_args_raw, fn_assignments, fn_call return fn_args_raw, fn_assignments, fn_call
mjx/mujoco/mjx/codegen/generate_warp_types.py
+1 -1
View File
@@ -440,7 +440,7 @@ def write_core_cls(
shape_property = _DATA_SHAPE_PROPERTY_FIELD shape_property = _DATA_SHAPE_PROPERTY_FIELD
new_class_body = _build_new_class_body_ast( new_class_body = _build_new_class_body_ast(
keys, keys, # pyrefly: ignore[bad-argument-type]
cls_name, cls_name,
annotations, annotations,
shape_property=shape_property, shape_property=shape_property,
mjx/mujoco/mjx/codegen/trace.py
+7 -7
View File
@@ -162,7 +162,7 @@ class _FunctionFieldUsageVisitor(ast.NodeVisitor):
called_fn_name = node.func.id called_fn_name = node.func.id
key = (hash(self._current_fpath), called_fn_name) key = (hash(self._current_fpath), called_fn_name)
if key not in self._visited_fns: if key not in self._visited_fns:
self._visited_fns.add(key) self._visited_fns.add(key) # pyrefly: ignore[bad-argument-type]
next_fpath = self._module_fpaths.get( next_fpath = self._module_fpaths.get(
called_fn_name, self._current_fpath called_fn_name, self._current_fpath
) )
@@ -181,8 +181,8 @@ class _FunctionFieldUsageVisitor(ast.NodeVisitor):
if len(node.args) != 2: if len(node.args) != 2:
raise ValueError(f'wp.copy() must have 2 arguments, got {node.args}.') raise ValueError(f'wp.copy() must have 2 arguments, got {node.args}.')
out_node, in_node = node.args out_node, in_node = node.args
self.add_field_usage(out_node, True) self.add_field_usage(out_node, True) # pyrefly: ignore[bad-argument-type]
self.add_field_usage(in_node, False) self.add_field_usage(in_node, False) # pyrefly: ignore[bad-argument-type]
return return
# do not trace array creation # do not trace array creation
@@ -201,7 +201,7 @@ class _FunctionFieldUsageVisitor(ast.NodeVisitor):
key = (hash(self._current_fpath), called_fn_name) key = (hash(self._current_fpath), called_fn_name)
next_fpath = self._module_fpaths.get(parts[0]) next_fpath = self._module_fpaths.get(parts[0])
if next_fpath and key not in self._visited_fns: if next_fpath and key not in self._visited_fns:
self._visited_fns.add(key) self._visited_fns.add(key) # pyrefly: ignore[bad-argument-type]
self.recurse_trace(next_fpath, called_fn_name) self.recurse_trace(next_fpath, called_fn_name)
self.generic_visit(node) self.generic_visit(node)
@@ -223,10 +223,10 @@ def trace_function(
) -> FieldUsage: ) -> FieldUsage:
"""Traces the function statically to find usages of model and data fields.""" """Traces the function statically to find usages of model and data fields."""
base_path = file.get_base_path() base_path = file.get_base_path()
fpath = base_path / fpath fpath = base_path / fpath # pyrefly: ignore[bad-assignment]
logging.info('Tracing function: "%s" in "%s"', fn, fpath) logging.info('Tracing function: "%s" in "%s"', fn, fpath)
src = fpath.read_text() src = fpath.read_text() # pyrefly: ignore[missing-attribute]
parsed_ast = ast.parse(src, filename=str(fpath)) parsed_ast = ast.parse(src, filename=str(fpath))
target_fn_nodes = ( target_fn_nodes = (
@@ -262,7 +262,7 @@ def trace_function(
if visited_fns is None: if visited_fns is None:
visited_fns = set() visited_fns = set()
visitor = _FunctionFieldUsageVisitor(fpath, visited_fns, mjwarp_field_info) visitor = _FunctionFieldUsageVisitor(fpath, visited_fns, mjwarp_field_info) # pyrefly: ignore[bad-argument-type]
for body in target_fn_node.body: for body in target_fn_node.body:
visitor.visit(body) visitor.visit(body)
mjx/mujoco/mjx/warp/ffi.py
+6 -6
View File
@@ -119,7 +119,7 @@ def jax_callable_variadic_tuple(
# Provide a flattened signature for the Warp callable machinery. # Provide a flattened signature for the Warp callable machinery.
new_signature = flatten_signature(inspect.signature(func), args) new_signature = flatten_signature(inspect.signature(func), args)
func_wrapper.__signature__ = new_signature func_wrapper.__signature__ = new_signature # pyrefly: ignore[missing-attribute]
func_wrapper.__annotations__ = { func_wrapper.__annotations__ = {
p.name: p.annotation p.name: p.annotation
for p in new_signature.parameters.values() for p in new_signature.parameters.values()
@@ -205,7 +205,7 @@ def _expand_dim_from_path(
if ndim is None or ndim < 0: if ndim is None or ndim < 0:
return leaf return leaf
if ndim > leaf.ndim: if ndim > leaf.ndim:
leaf = jp.expand_dims(leaf, axis=np.arange(ndim - leaf.ndim)) leaf = jp.expand_dims(leaf, axis=np.arange(ndim - leaf.ndim)) # pyrefly: ignore[bad-argument-type]
if ndim != leaf.ndim: if ndim != leaf.ndim:
raise AssertionError( raise AssertionError(
f'Leaf node ndim ({leaf.ndim}) and expected ndim ({ndim}) do not match' f'Leaf node ndim ({leaf.ndim}) and expected ndim ({ndim}) do not match'
@@ -221,7 +221,7 @@ def _squeeze_dim(leaf_expanded: Any, leaf: Any) -> Any:
f' ndim {leaf.ndim}' f' ndim {leaf.ndim}'
) )
if leaf_expanded.ndim > leaf.ndim: if leaf_expanded.ndim > leaf.ndim:
return jp.squeeze(leaf_expanded, np.arange(leaf_expanded.ndim - leaf.ndim)) return jp.squeeze(leaf_expanded, np.arange(leaf_expanded.ndim - leaf.ndim)) # pyrefly: ignore[bad-argument-type]
return leaf_expanded return leaf_expanded
@@ -288,7 +288,7 @@ def _maybe_broadcast_to(
"""Broadcasts fields that are used in MuJoCo Warp.""" """Broadcasts fields that are used in MuJoCo Warp."""
ndim = _get_mapping_from_tree_path(path, mjx_warp_types._NDIM[cls_str]) ndim = _get_mapping_from_tree_path(path, mjx_warp_types._NDIM[cls_str])
needs_batch_dim = _get_mapping_from_tree_path( needs_batch_dim = _get_mapping_from_tree_path(
path, mjx_warp_types._BATCH_DIM[cls_str] # pylint: disable=protected-access path, mjx_warp_types._BATCH_DIM[cls_str] # pylint: disable=protected-access # pyrefly: ignore[bad-argument-type]
) )
needs_batch_dim = bool(needs_batch_dim) and (ndim is not None and ndim > 0) needs_batch_dim = bool(needs_batch_dim) and (ndim is not None and ndim > 0)
if needs_batch_dim and not is_batched: if needs_batch_dim and not is_batched:
@@ -305,7 +305,7 @@ def _check_leading_dim(
): ):
"""Asserts that the batch dimension of a leaf node matches the expected batch dimension.""" """Asserts that the batch dimension of a leaf node matches the expected batch dimension."""
has_batch_dim = _get_mapping_from_tree_path( has_batch_dim = _get_mapping_from_tree_path(
path, mjx_warp_types._BATCH_DIM['Data'] path, mjx_warp_types._BATCH_DIM['Data'] # pyrefly: ignore[bad-argument-type]
) )
attr = tree_path_to_attr_str(path) attr = tree_path_to_attr_str(path)
if has_batch_dim and leaf.shape[0] != expected_batch_dim: if has_batch_dim and leaf.shape[0] != expected_batch_dim:
@@ -396,7 +396,7 @@ def marshal_custom_vmap(
out_batched = jax.tree.map_with_path( out_batched = jax.tree.map_with_path(
# NB: if a field is not in MuJoCo Warp, we let JAX do its magic. # NB: if a field is not in MuJoCo Warp, we let JAX do its magic.
lambda path, x: _get_mapping_from_tree_path( lambda path, x: _get_mapping_from_tree_path(
path, mjx_warp_types._BATCH_DIM['Data'] # pylint: disable=protected-access path, mjx_warp_types._BATCH_DIM['Data'] # pylint: disable=protected-access # pyrefly: ignore[bad-argument-type]
) )
or x, or x,
out_batched, out_batched,
python/mujoco/__init__.py
+2 -2
View File
@@ -232,8 +232,8 @@ def _bind_data(
_specs.MjSpec.from_zip = from_zip _specs.MjSpec.from_zip = from_zip
_specs.MjSpec.to_zip = to_zip _specs.MjSpec.to_zip = to_zip
_structs.MjData.bind = _bind_data _structs.MjData.bind = _bind_data # pyrefly: ignore[bad-assignment]
_structs.MjModel.bind = _bind_model _structs.MjModel.bind = _bind_model # pyrefly: ignore[bad-assignment]
HEADERS_DIR = os.path.join(os.path.dirname(__file__), 'include/mujoco') HEADERS_DIR = os.path.join(os.path.dirname(__file__), 'include/mujoco')
PLUGINS_DIR = os.path.join(os.path.dirname(__file__), 'plugin') PLUGINS_DIR = os.path.join(os.path.dirname(__file__), 'plugin')
python/mujoco/codegen/generate_spec_bindings.py
+6 -6
View File
@@ -142,10 +142,10 @@ def _struct_binding_code(
for f in field.fields for f in field.fields
): ):
for subfield in field.fields: for subfield in field.fields:
code += _binding_code(subfield, name) code += _binding_code(subfield, name) # pyrefly: ignore[bad-argument-type]
# generate for the struct itself # generate for the struct itself
field = ast_nodes.ValueType(name=name) field = ast_nodes.ValueType(name=name) # pyrefly: ignore[bad-assignment]
code += _value_binding_code(field, classname, varname) code += _value_binding_code(field, classname, varname) # pyrefly: ignore[bad-argument-type]
return code return code
@@ -211,7 +211,7 @@ def _ptr_binding_code(
if vartype == 'mjsElement': # this is ignored by the caller if vartype == 'mjsElement': # this is ignored by the caller
return 'mjsElement' return 'mjsElement'
if vartype.startswith('mjs'): # for structs, use the value case if vartype.startswith('mjs'): # for structs, use the value case
return _value_binding_code(field.inner_type, classname, varname) return _value_binding_code(field.inner_type, classname, varname) # pyrefly: ignore[bad-argument-type]
elif vartype == 'mjString': # C++ string -> Python string elif vartype == 'mjString': # C++ string -> Python string
return f"""\ return f"""\
{classname}.def_property( {classname}.def_property(
@@ -347,7 +347,7 @@ def generate() -> None:
) and key != 'mjsElement': ) and key != 'mjsElement':
print('\n // ' + key) print('\n // ' + key)
for field in structs.STRUCTS[key].fields: for field in structs.STRUCTS[key].fields:
code = _binding_code(field, key) code = _binding_code(field, key) # pyrefly: ignore[bad-argument-type]
if code != 'mjsElement': if code != 'mjsElement':
print(code) print(code)
@@ -556,7 +556,7 @@ def generate_add() -> None:
py_args = ['py::arg("name") = py::none()'] py_args = ['py::arg("name") = py::none()']
for field in structs.STRUCTS[key].fields: for field in structs.STRUCTS[key].fields:
line, set_type, name, type_name, cpp_arg, py_arg = _field(field) line, set_type, name, type_name, cpp_arg, py_arg = _field(field) # pyrefly: ignore[bad-argument-type]
if line: if line:
code_field = code_field + '\n ' + line code_field = code_field + '\n ' + line
set_types.append(set_type) set_types.append(set_type)
python/mujoco/introspect/codegen/generate.py
+4 -4
View File
@@ -126,22 +126,22 @@ def main(argv: Sequence[str]) -> None:
raise app.UsageError('Too many command-line arguments.') raise app.UsageError('Too many command-line arguments.')
visitor = ast_processor.process( visitor = ast_processor.process(
_JSON_PATH.value, _JSON_PATH.value, # pyrefly: ignore[bad-argument-type]
(_HEADER_PATHS.value or '').split(), (_HEADER_PATHS.value or '').split(),
_EXCLUDED, _EXCLUDED,
) )
with open(_OUT_FUNCTIONS.value, 'w') as f: with open(_OUT_FUNCTIONS.value, 'w') as f: # pyrefly: ignore[no-matching-overload]
functions_str = formatter.format_as_python_code(visitor.exported_functions) functions_str = formatter.format_as_python_code(visitor.exported_functions)
f.write(_HEADER_TEMPLATE.format(year=2022, type='functions')) f.write(_HEADER_TEMPLATE.format(year=2022, type='functions'))
f.write(_FUNCTIONS_TEMPLATE.format(functions_str=functions_str)) f.write(_FUNCTIONS_TEMPLATE.format(functions_str=functions_str))
with open(_OUT_ENUMS.value, 'w') as f: with open(_OUT_ENUMS.value, 'w') as f: # pyrefly: ignore[no-matching-overload]
enums_str = formatter.format_as_python_code(visitor.exported_enums) enums_str = formatter.format_as_python_code(visitor.exported_enums)
f.write(_HEADER_TEMPLATE.format(year=2022, type='enums')) f.write(_HEADER_TEMPLATE.format(year=2022, type='enums'))
f.write(_ENUMS_TEMPLATE.format(enums_str=enums_str)) f.write(_ENUMS_TEMPLATE.format(enums_str=enums_str))
with open(_OUT_STRUCTS.value, 'w') as f: with open(_OUT_STRUCTS.value, 'w') as f: # pyrefly: ignore[no-matching-overload]
structs_str = formatter.format_as_python_code(visitor.exported_structs) structs_str = formatter.format_as_python_code(visitor.exported_structs)
f.write(_HEADER_TEMPLATE.format(year=2023, type='structs')) f.write(_HEADER_TEMPLATE.format(year=2023, type='structs'))
f.write(_STRUCTS_TEMPLATE.format(structs_str=structs_str)) f.write(_STRUCTS_TEMPLATE.format(structs_str=structs_str))
python/mujoco/introspect/type_parsing.py
+2 -2
View File
@@ -57,7 +57,7 @@ def _parse_maybe_array(
int(s.strip()) for s in ARRAY_N_PATTERN.findall(array_match.group(0))) int(s.strip()) for s in ARRAY_N_PATTERN.findall(array_match.group(0)))
inner_type_str = type_name[:array_match.start()] inner_type_str = type_name[:array_match.start()]
return ast_nodes.ArrayType( return ast_nodes.ArrayType(
inner_type=_parse_maybe_pointer(inner_type_str.strip(), innermost_type), inner_type=_parse_maybe_pointer(inner_type_str.strip(), innermost_type), # pyrefly: ignore[bad-argument-type]
extents=extents) extents=extents)
else: else:
return _parse_maybe_pointer(type_name, innermost_type) return _parse_maybe_pointer(type_name, innermost_type)
@@ -140,7 +140,7 @@ def parse_type(
result = None result = None
while type_str_stack: while type_str_stack:
try: try:
result = _parse_maybe_array(type_str_stack.pop(), result) result = _parse_maybe_array(type_str_stack.pop(), result) # pyrefly: ignore[bad-argument-type]
except AssertionError as e: except AssertionError as e:
raise ValueError(f'invalid type name {type_name!r}') from e raise ValueError(f'invalid type name {type_name!r}') from e
python/mujoco/minimize.py
+1 -1
View File
@@ -405,7 +405,7 @@ def least_squares(
print(message, file=output) print(message, file=output)
# Append log to trace, call iter_callback. # Append log to trace, call iter_callback.
log = IterLog(candidate=x, objective=y, reduction=reduction, regularizer=mu) log = IterLog(candidate=x, objective=y, reduction=reduction, regularizer=mu) # pyrefly: ignore[bad-argument-type]
if verbose >= Verbosity.FULLITER.value: if verbose >= Verbosity.FULLITER.value:
log = dataclasses.replace( log = dataclasses.replace(
log, residual=r, jacobian=jac, grad=grad / D, step=step log, residual=r, jacobian=jac, grad=grad / D, step=step
python/mujoco/rendering/classic/renderer.py
+3 -3
View File
@@ -93,7 +93,7 @@ the clause:
mujoco.mjr_setBuffer( mujoco.mjr_setBuffer(
mujoco.mjtFramebuffer.mjFB_OFFSCREEN.value, self._mjr_context mujoco.mjtFramebuffer.mjFB_OFFSCREEN.value, self._mjr_context
) )
self._mjr_context.readDepthMap = mujoco.mjtDepthMap.mjDEPTH_ZEROFAR self._mjr_context.readDepthMap = mujoco.mjtDepthMap.mjDEPTH_ZEROFAR # pyrefly: ignore[bad-assignment]
# Default render flags. # Default render flags.
self._depth_rendering = False self._depth_rendering = False
@@ -291,10 +291,10 @@ the clause:
# Defaults to mjCAMERA_FREE, otherwise mjCAMERA_FIXED refers to a # Defaults to mjCAMERA_FREE, otherwise mjCAMERA_FIXED refers to a
# camera explicitly defined in the model. # camera explicitly defined in the model.
if camera_id == -1: if camera_id == -1:
camera.type = mujoco.mjtCamera.mjCAMERA_FREE camera.type = mujoco.mjtCamera.mjCAMERA_FREE # pyrefly: ignore[bad-assignment]
mujoco.mjv_defaultFreeCamera(self._model, camera) mujoco.mjv_defaultFreeCamera(self._model, camera)
else: else:
camera.type = mujoco.mjtCamera.mjCAMERA_FIXED camera.type = mujoco.mjtCamera.mjCAMERA_FIXED # pyrefly: ignore[bad-assignment]
scene_option = scene_option or self._scene_option scene_option = scene_option or self._scene_option
mujoco.mjv_updateScene( mujoco.mjv_updateScene(
python/mujoco/rollout.py
+3 -3
View File
@@ -180,7 +180,7 @@ class Rollout:
model = [model] # Use a length 1 list to simplify code below model = [model] # Use a length 1 list to simplify code below
if not isinstance(data, list): if not isinstance(data, list):
data = [data] # Use a length 1 list to simplify code below data = [data] # Use a length 1 list to simplify code below # pyrefly: ignore[bad-assignment]
# infer nstep, check for incompatibilities # infer nstep, check for incompatibilities
nstep = _infer_dimension( nstep = _infer_dimension(
@@ -310,9 +310,9 @@ def rollout(
ValueError: bad shapes or sizes. ValueError: bad shapes or sizes.
""" # fmt: skip """ # fmt: skip
if not isinstance(data, list): if not isinstance(data, list):
data = [data] # Use a length 1 list to simplify code below data = [data] # Use a length 1 list to simplify code below # pyrefly: ignore[bad-assignment]
nthread = len(data) if len(data) > 1 else 0 nthread = len(data) if len(data) > 1 else 0 # pyrefly: ignore[bad-argument-type]
# Use a persistent thread pool if requested # Use a persistent thread pool if requested
if persistent_pool: if persistent_pool:
python/mujoco/sysid/_src/optimize.py
+1 -1
View File
@@ -53,7 +53,7 @@ def _warn_if_ill_conditioned(
residual=f, residual=f,
x=x0.reshape(-1, 1), x=x0.reshape(-1, 1),
r=r0, r=r0,
eps=eps, eps=eps, # pyrefly: ignore[bad-argument-type]
n_res=0, n_res=0,
bounds=[bounds[0].reshape(-1, 1), bounds[1].reshape(-1, 1)], bounds=[bounds[0].reshape(-1, 1), bounds[1].reshape(-1, 1)],
)[0], )[0],
python/mujoco/viewer.py
+2 -2
View File
@@ -325,8 +325,8 @@ def _reload(
def _physics_loop(simulate: _Simulate, loader: Optional[_InternalLoaderType]): def _physics_loop(simulate: _Simulate, loader: Optional[_InternalLoaderType]):
"""Physics loop for the GUI, to be run in a separate thread.""" """Physics loop for the GUI, to be run in a separate thread."""
m: mujoco.MjModel = None m: mujoco.MjModel = None # pyrefly: ignore[bad-assignment]
d: mujoco.MjData = None d: mujoco.MjData = None # pyrefly: ignore[bad-assignment]
ctrl_noise = np.array([]) ctrl_noise = np.array([])
reload = True reload = True