diff --git a/mjx/mujoco/mjx/_src/collision_convex.py b/mjx/mujoco/mjx/_src/collision_convex.py
index d4662a09..61950b68 100644
--- a/mjx/mujoco/mjx/_src/collision_convex.py
+++ b/mjx/mujoco/mjx/_src/collision_convex.py
@@ -56,27 +56,27 @@ def collider(ncon: int):
       fn = collision_fn
       for i in [0, 1]:
         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(
               pos=0, mat=0, size=0, face=0, vert=0
           )
         elif key.types[i] == GeomType.MESH:
           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(
               pos=0, mat=0, size=0
           )
         elif key.types[i] == GeomType.HFIELD:
           hfield_info = mesh.hfield(m, key.data_ids[i])
-          infos[i] = hfield_info.replace(pos=infos[i].pos, mat=infos[i].mat)
-          in_axes[i] = hfield_info.replace(pos=0, mat=0, data=None)
+          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)  # pyrefly: ignore[unsupported-operation]
           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
       if ncon > 1:
         return jax.tree_util.tree_map(jp.concatenate, (dist, pos, frame))
       return dist, pos, frame
 
-    collide.ncon = ncon
+    collide.ncon = ncon  # pyrefly: ignore[missing-attribute]
     return collide
 
   return wrapper
@@ -1051,7 +1051,7 @@ def hfield_sphere(
 ) -> Collision:
   """Calculates contacts between a hfield and a sphere."""
   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)
   mask = dist < jp.minimum(0, dist.min() + 1e-3)
@@ -1075,7 +1075,7 @@ def hfield_capsule(
 ) -> Collision:
   """Calculates contacts between a hfield and a capsule."""
   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)
   mask = dist < jp.minimum(0, dist.min() + 1e-3)
@@ -1099,7 +1099,7 @@ def hfield_convex(
 ) -> Collision:
   """Calculates contacts between a hfield and a capsule."""
   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)
   mask = dist < jp.minimum(0, dist.min() + 1e-3)
diff --git a/mjx/mujoco/mjx/_src/collision_driver.py b/mjx/mujoco/mjx/_src/collision_driver.py
index 63e8c6e9..abf5d227 100644
--- a/mjx/mujoco/mjx/_src/collision_driver.py
+++ b/mjx/mujoco/mjx/_src/collision_driver.py
@@ -384,7 +384,7 @@ def make_condim(
     if k.types[1] == mujoco.mjtGeom.mjGEOM_SDF:
       ncon = sdf_initpoints
     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:
         ncon = func.ncon  # pytype: disable=attribute-error
       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)
 
     # 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
 
     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)], [])
   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]
diff --git a/mjx/mujoco/mjx/_src/collision_primitive.py b/mjx/mujoco/mjx/_src/collision_primitive.py
index b7d695a9..5a04a93a 100644
--- a/mjx/mujoco/mjx/_src/collision_primitive.py
+++ b/mjx/mujoco/mjx/_src/collision_primitive.py
@@ -40,7 +40,7 @@ def collider(ncon: int):
         return jax.tree_util.tree_map(jp.concatenate, (dist, pos, frame))
       return dist, pos, frame
 
-    collide.ncon = ncon
+    collide.ncon = ncon  # pyrefly: ignore[missing-attribute]
     return collide
 
   return wrapper
diff --git a/mjx/mujoco/mjx/_src/collision_sdf.py b/mjx/mujoco/mjx/_src/collision_sdf.py
index 3bf8b824..d5754381 100644
--- a/mjx/mujoco/mjx/_src/collision_sdf.py
+++ b/mjx/mujoco/mjx/_src/collision_sdf.py
@@ -54,7 +54,7 @@ def collider(ncon: int):
         return jax.tree_util.tree_map(jp.concatenate, (dist, pos, frame))
       return dist, pos, frame
 
-    collide.ncon = ncon
+    collide.ncon = ncon  # pyrefly: ignore[missing-attribute]
     return collide
 
   return wrapper
diff --git a/mjx/mujoco/mjx/_src/constraint.py b/mjx/mujoco/mjx/_src/constraint.py
index 042d2b72..4c5bc298 100644
--- a/mjx/mujoco/mjx/_src/constraint.py
+++ b/mjx/mujoco/mjx/_src/constraint.py
@@ -716,8 +716,8 @@ def make_constraint(m: Model, d: Data) -> Data:
 
   if not efcs:
     z = jp.empty(0)
-    d = d.tree_replace({'_impl.efc_J': jp.empty((0, m.nv))})
-    d = d.tree_replace({
+    d = d.tree_replace({'_impl.efc_J': jp.empty((0, m.nv))})  # pyrefly: ignore[bad-assignment]
+    d = d.tree_replace({  # pyrefly: ignore[bad-assignment]
         '_impl.efc_D': z,
         '_impl.efc_aref': 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
 
   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_D': 1 / r,
       '_impl.efc_aref': aref,
       '_impl.efc_pos': pos,
       '_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
diff --git a/mjx/mujoco/mjx/_src/dataclasses.py b/mjx/mujoco/mjx/_src/dataclasses.py
index fad9ce7c..f8c37403 100644
--- a/mjx/mujoco/mjx/_src/dataclasses.py
+++ b/mjx/mujoco/mjx/_src/dataclasses.py
@@ -75,7 +75,7 @@ def dataclass(clz: _T, register_as_pytree: bool) -> _T:
   Returns:
     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
 
   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
     )
 
-  return data_clz
+  return data_clz  # pyrefly: ignore[bad-return]
 
 
 TNode = TypeVar('TNode', bound='PyTreeNode')
@@ -163,7 +163,7 @@ class PyTreeNode:
 
   @classmethod
   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(
       self, params: Dict[str, Optional[jax.typing.ArrayLike]]
@@ -190,7 +190,7 @@ def _tree_replace(
     for i, g in enumerate(lst):
       if not hasattr(g, attr[1]):
         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)
 
     return base.replace(**{attr[0]: lst})
diff --git a/mjx/mujoco/mjx/_src/forward.py b/mjx/mujoco/mjx/_src/forward.py
index cc92a829..46257215 100644
--- a/mjx/mujoco/mjx/_src/forward.py
+++ b/mjx/mujoco/mjx/_src/forward.py
@@ -89,7 +89,7 @@ def fwd_velocity(m: Model, d: Data) -> Data:
   if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
     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.ten_velocity': d._impl.ten_J @ d.qvel,
   })
@@ -357,7 +357,7 @@ def euler(m: Model, d: Data) -> Data:
     else:
       M = d._impl.M + jp.diag(m.opt.timestep * m.dof_damping)
     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
     qacc = smooth.solve_m(m, dh, qfrc)
   return _advance(m, d, d.act_dot, qacc)
diff --git a/mjx/mujoco/mjx/_src/inverse.py b/mjx/mujoco/mjx/_src/inverse.py
index 2b045ea4..18055420 100644
--- a/mjx/mujoco/mjx/_src/inverse.py
+++ b/mjx/mujoco/mjx/_src/inverse.py
@@ -76,7 +76,7 @@ def inv_constraint(m: Model, d: Data) -> Data:
   # update
   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,
       '_impl.efc_force': ctx.efc_force,
   })
diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py
index 9f987ee2..09165836 100644
--- a/mjx/mujoco/mjx/_src/io.py
+++ b/mjx/mujoco/mjx/_src/io.py
@@ -151,15 +151,15 @@ def _resolve_impl_and_device(
   if (has_impl, has_device) == (True, True):
     pass
   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):
-    impl = _resolve_impl(device)
+    impl = _resolve_impl(device)  # pyrefly: ignore[bad-argument-type]
   else:
     device = jax.devices()[0]
     logging.info('Using JAX default device: %s.', 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
 
 
@@ -724,7 +724,7 @@ def _make_data_jax(
       qpos=jp.array(m.qpos0, dtype=float_),
       eq_active=m.eq_active0,
       _impl=impl,
-      **_make_data_public_fields(m),
+      **_make_data_public_fields(m),  # pyrefly: ignore[bad-argument-type]
   )
 
   if m.nmocap:
@@ -773,7 +773,7 @@ def _make_data_warp(
         nvmax=nvmax,
     )  # 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:
     if k in {'userdata', 'plugin_state', 'history'}:
       continue
@@ -1287,14 +1287,14 @@ def _get_data_into_warp(
         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.nacon[0]
     nefc = int(d_i._impl.nefc)
     # nj = int(d_i._impl.nj[0])
     nj = 0  # TODO(btaba): add nj back
 
-    if ncon != result_i.ncon or nefc != result_i.nefc or nj != result_i.nJ:
-      mujoco._functions._realloc_con_efc(result_i, ncon=ncon, nefc=nefc, nJ=nj)  # pylint: disable=protected-access
+    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  # pyrefly: ignore[bad-argument-type]
 
     all_fields = types.Data.fields() + mjxw.types.DataWarp.fields()
     for field in all_fields:
@@ -1368,14 +1368,14 @@ def _get_data_into(
 
   for i in range(batch_size):
     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()
     efc_active = (d_i._impl.efc_J != 0).any(axis=1)
     nefc = int(efc_active.sum())
     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:
-      mujoco._functions._realloc_con_efc(result_i, ncon=ncon, nefc=nefc, nJ=nj)  # pylint: disable=protected-access
+    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  # pyrefly: ignore[bad-argument-type]
 
     if d.impl == types.Impl.JAX:
       all_fields = types.Data.fields() + types.DataJAX.fields()
@@ -1389,10 +1389,10 @@ def _get_data_into(
         continue
 
       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_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
 
       # MuJoCo actuator_moment is sparse, MJX uses a dense representation.
@@ -1412,10 +1412,10 @@ def _get_data_into(
             )
           else:
             actuator_moment = d_i._impl.actuator_moment
-        result_i.moment_rownnz[:] = moment_rownnz
-        result_i.moment_rowadr[:] = moment_rowadr
-        result_i.moment_colind[:] = moment_colind
-        result_i.actuator_moment[:] = actuator_moment
+        result_i.moment_rownnz[:] = moment_rownnz  # pyrefly: ignore[missing-attribute]
+        result_i.moment_rowadr[:] = moment_rowadr  # pyrefly: ignore[missing-attribute]
+        result_i.moment_colind[:] = moment_colind  # pyrefly: ignore[missing-attribute]
+        result_i.actuator_moment[:] = actuator_moment  # pyrefly: ignore[missing-attribute]
         continue
 
       # MuJoCo ten_J is sparse, MJX uses a dense representation.
@@ -1435,7 +1435,7 @@ def _get_data_into(
             )
           else:
             ten_j = d_i._impl.ten_J
-        result_i.ten_J[:] = ten_j
+        result_i.ten_J[:] = ten_j  # pyrefly: ignore[missing-attribute]
         continue
 
       if hasattr(d_i._impl, field.name):
@@ -1461,9 +1461,9 @@ def _get_data_into(
               efc_J_rowadr,
               efc_J_colind,
           )
-          result_i.efc_J_rownnz[:] = efc_J_rownnz
-          result_i.efc_J_rowadr[:] = efc_J_rowadr
-          result_i.efc_J_colind[:] = efc_J_colind
+          result_i.efc_J_rownnz[:] = efc_J_rownnz  # pyrefly: ignore[missing-attribute]
+          result_i.efc_J_rowadr[:] = efc_J_rowadr  # pyrefly: ignore[missing-attribute]
+          result_i.efc_J_colind[:] = efc_J_colind  # pyrefly: ignore[missing-attribute]
           value = efc_J
         else:
           value = value.reshape(-1)
@@ -1496,7 +1496,7 @@ def _get_data_into(
 
     # recalculate qLD and qLDiagInv as MJX and MuJoCo have different
     # 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
@@ -1533,7 +1533,7 @@ def _get_data_into_cpp(
     d_i: types.Data = (
         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:
       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
       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(
@@ -1810,5 +1810,5 @@ def create_render_context(
   _check_warp_installed()
   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(
-      mjm, nworld=nworld, devices=devices, **kwargs
+      mjm, nworld=nworld, devices=devices, **kwargs  # pyrefly: ignore[bad-argument-type]
   )
diff --git a/mjx/mujoco/mjx/_src/ray.py b/mjx/mujoco/mjx/_src/ray.py
index 5d4b5c3e..aaafc5a3 100644
--- a/mjx/mujoco/mjx/_src/ray.py
+++ b/mjx/mujoco/mjx/_src/ray.py
@@ -263,8 +263,8 @@ def ray(
   for bodyid in bodyexclude:
     geom_filter &= (m.geom_bodyid != bodyid)
   if geomgroup:
-    geomgroup = np.array(geomgroup, dtype=bool)
-    geom_filter &= geomgroup[np.clip(m.geom_group, 0, mujoco.mjNGROUP)]
+    geomgroup = np.array(geomgroup, dtype=bool)  # pyrefly: ignore[bad-assignment]
+    geom_filter &= geomgroup[np.clip(m.geom_group, 0, mujoco.mjNGROUP)]  # pyrefly: ignore[bad-index]
 
   # map ray to local geom frames
   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_]
 
     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:
       dist = jax.vmap(fn)(*args)
 
@@ -314,4 +314,4 @@ def ray_geom(
   Returns:
     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]
diff --git a/mjx/mujoco/mjx/_src/scan.py b/mjx/mujoco/mjx/_src/scan.py
index 0bcfc195..7bfac7d9 100644
--- a/mjx/mujoco/mjx/_src/scan.py
+++ b/mjx/mujoco/mjx/_src/scan.py
@@ -331,7 +331,7 @@ def flat(
     reordered_ys.append(_take(y, _index(ids, input_order)))
   y = reordered_ys if f_ret_is_seq else reordered_ys[0]
 
-  return y
+  return y  # pyrefly: ignore[bad-return]
 
 
 def body_tree(
@@ -443,7 +443,7 @@ def body_tree(
       if i < len(in_types):
         key_in_take.setdefault(key, []).append(ids)
       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)
   keys = sorted(key_body_ids, reverse=reverse)
@@ -494,4 +494,4 @@ def body_tree(
 
   y = y[0] if len(out_types) == 1 else y
 
-  return y
+  return y  # pyrefly: ignore[bad-return]
diff --git a/mjx/mujoco/mjx/_src/sensor.py b/mjx/mujoco/mjx/_src/sensor.py
index 7e0930cd..d3ae2da8 100644
--- a/mjx/mujoco/mjx/_src/sensor.py
+++ b/mjx/mujoco/mjx/_src/sensor.py
@@ -339,9 +339,9 @@ def sensor_vel(m: Model, d: Data) -> Data:
     elif sensor_type == SensorType.JOINTVEL:
       sensor = d.qvel[m.jnt_dofadr[objid]]
     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:
-      sensor = d._impl.actuator_velocity[objid]
+      sensor = d._impl.actuator_velocity[objid]  # pyrefly: ignore[missing-attribute]
     elif sensor_type == SensorType.BALLANGVEL:
       jnt_dotadr = m.jnt_dofadr[objid, None] + np.arange(3)[None]
       sensor = d.qvel[jnt_dotadr]
@@ -403,10 +403,10 @@ def sensor_vel(m: Model, d: Data) -> Data:
         adrs.append(adrt.reshape(-1))
       continue  # avoid adding to sensors/adrs list a second time
     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)
     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)
     else:
       # 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
       conray = jax.vmap(
           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)
 
       # 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
         force_mag = jax.vmap(
             lambda forcetorque: jp.dot(forcetorque[:3], forcetorque[:3])
-        )(contact_force)
+        )(contact_force)  # pyrefly: ignore[unbound-name]
 
       def _reduce(reduction, mask):
         if reduction == 1:  # mindist
@@ -662,7 +662,7 @@ def sensor_acc(m: Model, d: Data) -> Data:
           slot.append(jp.repeat(nfound, num)[:, None])
 
         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
           slot.append(flip * contact_force[cid, 3:])
@@ -699,14 +699,14 @@ def sensor_acc(m: Model, d: Data) -> Data:
       bodyid = m.site_bodyid[objid]
       rot = d.site_xmat[objid]
       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]]
 
       sensor = _accelerometer(cvel, cacc, dif, rot)
       adr = (adr[:, None] + np.arange(3)[None]).reshape(-1)
     elif sensor_type == SensorType.FORCE:
       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]
       sensor = jax.vmap(lambda mat, vec: mat.T @ vec)(
           site_xmat, cfrc_int[:, 3:]
@@ -715,7 +715,7 @@ def sensor_acc(m: Model, d: Data) -> Data:
     elif sensor_type == SensorType.TORQUE:
       bodyid = m.site_bodyid[objid]
       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]
       dif = d.site_xpos[objid] - d.subtree_com[rootid]
       sensor = jax.vmap(
@@ -744,7 +744,7 @@ def sensor_acc(m: Model, d: Data) -> Data:
         pos, bodyid = objtype_data[ot]
         pos = pos[objidt]
         bodyid = bodyid[objidt]
-        cacc = d._impl.cacc[bodyid]
+        cacc = d._impl.cacc[bodyid]  # pyrefly: ignore[missing-attribute]
 
         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
       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)
 
   if not adrs:
diff --git a/mjx/mujoco/mjx/_src/smooth.py b/mjx/mujoco/mjx/_src/smooth.py
index 1e8b4df6..708fd84d 100644
--- a/mjx/mujoco/mjx/_src/smooth.py
+++ b/mjx/mujoco/mjx/_src/smooth.py
@@ -173,7 +173,7 @@ def com_pos(m: Model, d: Data) -> Data:
   root_com = subtree_com[m.body_rootid]
   offset = d.xipos - root_com
   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
   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.xaxis,
   )
-  d = d.tree_replace({'cdof': cdof})
+  d = d.tree_replace({'cdof': cdof})  # pyrefly: ignore[bad-assignment]
 
   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 = 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_cdof = jax.vmap(math.inert_mul)(crb_dof, d.cdof)
   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
 
 
@@ -318,7 +318,7 @@ def factor_m(m: Model, d: Data) -> Data:
 
   if not support.is_sparse(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
 
   depth = []
@@ -367,7 +367,7 @@ def factor_m(m: Model, d: Data) -> Data:
   qld_diag = qld[diag_adr]
   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
 
 
@@ -448,7 +448,7 @@ def com_vel(m: Model, d: Data) -> Data:
       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
 
@@ -544,7 +544,7 @@ def subtree_vel(m: Model, d: Data) -> Data:
       reverse=True,
   )
 
-  return d.tree_replace({
+  return d.tree_replace({  # pyrefly: ignore[bad-return]
       '_impl.subtree_linvel': subtree_linvel,
       '_impl.subtree_angmom': subtree_angmom,
   })
@@ -841,7 +841,7 @@ def rne_postconstraint(m: Model, d: Data) -> Data:
   )
 
   # update data
-  return d.tree_replace({
+  return d.tree_replace({  # pyrefly: ignore[bad-return]
       '_impl.cacc': cacc,
       '_impl.cfrc_int': cfrc_int,
       '_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)]
   ).reshape((m.nwrap, 2))
 
-  return d.tree_replace({
+  return d.tree_replace({  # pyrefly: ignore[bad-return]
       'ten_length': ten_length,
       '_impl.ten_J': ten_moment,
       '_impl.ten_wrapadr': jp.array(ten_wrapadr, dtype=int),
@@ -1276,7 +1276,7 @@ def transmission(m: Model, d: Data) -> Data:
       moment = jac @ wrench
     elif trntype == TrnType.TENDON:
       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:
       raise RuntimeError(f'unrecognized trntype: {TrnType(trntype)}')
 
@@ -1308,7 +1308,7 @@ def transmission(m: Model, d: Data) -> Data:
   length = length.reshape((m.nu,))
   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}
   )
   return d
@@ -1332,7 +1332,7 @@ def tendon_armature(m: Model, d: Data) -> Data:
     j = m.M_colind
     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:
@@ -1457,7 +1457,7 @@ def tendon_bias(m: Model, d: Data) -> Data:
   # add bias term: qfrc += ten_J * armature * ten_Jdot @ 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': (
           d.qfrc_bias
           + jp.sum(jax.vmap(jp.multiply)(d._impl.ten_J, coef), axis=0)
diff --git a/mjx/mujoco/mjx/_src/solver.py b/mjx/mujoco/mjx/_src/solver.py
index 6ac37f19..5bbd0220 100644
--- a/mjx/mujoco/mjx/_src/solver.py
+++ b/mjx/mujoco/mjx/_src/solver.py
@@ -482,7 +482,7 @@ def _linesearch(m: Model, d: Data, ctx: Context) -> Context:
     vv = jp.sum(v[:, 1:] * v[:, 1:], axis=1)
 
   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:
@@ -601,7 +601,7 @@ def solve(m: Model, d: Data) -> Data:
   else:
     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,
       'qacc': ctx.qacc,
       '_impl.efc_force': ctx.efc_force,
diff --git a/mjx/mujoco/mjx/_src/support.py b/mjx/mujoco/mjx/_src/support.py
index 914a0ba9..378a8054 100644
--- a/mjx/mujoco/mjx/_src/support.py
+++ b/mjx/mujoco/mjx/_src/support.py
@@ -501,15 +501,15 @@ class BindData(object):
         num = sum((typ == jt) * jt.dof_width() for jt in JointType)
       if isinstance(self.id, list):
         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))
         return self._slice(self.__getname(name), idx)
       elif num > 1:
         return self._slice(self.__getname(name), slice(adr, adr + num))
       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'):
-      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)
 
   def set(self, name: str, value: jax.Array) -> Data:
@@ -521,7 +521,7 @@ class BindData(object):
     try:
       iter(value)
     except TypeError:
-      value = [value]
+      value = [value]  # pyrefly: ignore[bad-assignment]
     if name in ('qpos', 'qvel', 'qacc', 'mocap_pos', 'mocap_quat'):
       adr = num = 0
       if name == 'qpos':
@@ -549,7 +549,7 @@ class BindData(object):
       num = [dim]
     i = 0
     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
       array = array.flatten().at[a : a + n].set(value[i : i + n]).reshape(shape)
       i += n
@@ -623,7 +623,7 @@ def contact_force_dim(
   if m.opt.cone == ConeType.PYRAMIDAL:
     efc_address = (
         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
     force = jax.vmap(_decode_pyramid, in_axes=(0, 0, None))(
diff --git a/mjx/mujoco/mjx/_src/test_util.py b/mjx/mujoco/mjx/_src/test_util.py
index 142a2eaa..f3157791 100644
--- a/mjx/mujoco/mjx/_src/test_util.py
+++ b/mjx/mujoco/mjx/_src/test_util.py
@@ -63,13 +63,13 @@ def benchmark(
   xla_flags += ' --xla_gpu_triton_gemm_any=True'
   os.environ['XLA_FLAGS'] = xla_flags
 
-  m.opt.solver = {
+  m.opt.solver = {  # pyrefly: ignore[bad-assignment]
       'cg': mujoco.mjtSolver.mjSOL_CG,
       'newton': mujoco.mjtSolver.mjSOL_NEWTON,
   }[solver.lower()]
   m.opt.iterations = iterations
   m.opt.ls_iterations = ls_iterations
-  m = io.put_model(m)
+  m = io.put_model(m)  # pyrefly: ignore[bad-assignment]
 
   @jax.pmap
   def init(key):
@@ -449,7 +449,7 @@ def create_mjcf(
     if p(80):
       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:
       geom1, geom2 = geom2, geom1
 
@@ -467,7 +467,7 @@ def create_mjcf(
     if p(50):
       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})
 
   # ElementTree.indent is not available before Python 3.9
diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py
index c6d422d5..50631dbe 100644
--- a/mjx/mujoco/mjx/_src/types.py
+++ b/mjx/mujoco/mjx/_src/types.py
@@ -105,10 +105,10 @@ class JointType(enum.IntEnum):
     HINGE: rotation angle (rad) around body-fixed axis  (1,)
   """
 
-  FREE = mujoco.mjtJoint.mjJNT_FREE
-  BALL = mujoco.mjtJoint.mjJNT_BALL
-  SLIDE = mujoco.mjtJoint.mjJNT_SLIDE
-  HINGE = mujoco.mjtJoint.mjJNT_HINGE
+  FREE = mujoco.mjtJoint.mjJNT_FREE  # pyrefly: ignore[bad-assignment]
+  BALL = mujoco.mjtJoint.mjJNT_BALL  # pyrefly: ignore[bad-assignment]
+  SLIDE = mujoco.mjtJoint.mjJNT_SLIDE  # pyrefly: ignore[bad-assignment]
+  HINGE = mujoco.mjtJoint.mjJNT_HINGE  # pyrefly: ignore[bad-assignment]
 
   def dof_width(self) -> int:
     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
   """
 
-  EULER = mujoco.mjtIntegrator.mjINT_EULER
-  RK4 = mujoco.mjtIntegrator.mjINT_RK4
-  IMPLICITFAST = mujoco.mjtIntegrator.mjINT_IMPLICITFAST
+  EULER = mujoco.mjtIntegrator.mjINT_EULER  # pyrefly: ignore[bad-assignment]
+  RK4 = mujoco.mjtIntegrator.mjINT_RK4  # pyrefly: ignore[bad-assignment]
+  IMPLICITFAST = mujoco.mjtIntegrator.mjINT_IMPLICITFAST  # pyrefly: ignore[bad-assignment]
   # unsupported: IMPLICIT
 
 
@@ -147,14 +147,14 @@ class GeomType(enum.IntEnum):
     SDF: signed distance field
   """
 
-  PLANE = mujoco.mjtGeom.mjGEOM_PLANE
-  HFIELD = mujoco.mjtGeom.mjGEOM_HFIELD
-  SPHERE = mujoco.mjtGeom.mjGEOM_SPHERE
-  CAPSULE = mujoco.mjtGeom.mjGEOM_CAPSULE
-  ELLIPSOID = mujoco.mjtGeom.mjGEOM_ELLIPSOID
-  CYLINDER = mujoco.mjtGeom.mjGEOM_CYLINDER
-  BOX = mujoco.mjtGeom.mjGEOM_BOX
-  MESH = mujoco.mjtGeom.mjGEOM_MESH
+  PLANE = mujoco.mjtGeom.mjGEOM_PLANE  # pyrefly: ignore[bad-assignment]
+  HFIELD = mujoco.mjtGeom.mjGEOM_HFIELD  # pyrefly: ignore[bad-assignment]
+  SPHERE = mujoco.mjtGeom.mjGEOM_SPHERE  # pyrefly: ignore[bad-assignment]
+  CAPSULE = mujoco.mjtGeom.mjGEOM_CAPSULE  # pyrefly: ignore[bad-assignment]
+  ELLIPSOID = mujoco.mjtGeom.mjGEOM_ELLIPSOID  # pyrefly: ignore[bad-assignment]
+  CYLINDER = mujoco.mjtGeom.mjGEOM_CYLINDER  # pyrefly: ignore[bad-assignment]
+  BOX = mujoco.mjtGeom.mjGEOM_BOX  # pyrefly: ignore[bad-assignment]
+  MESH = mujoco.mjtGeom.mjGEOM_MESH  # pyrefly: ignore[bad-assignment]
   # unsupported: NGEOMTYPES, ARROW*, LINE, SKIN, LABEL, NONE
 
 
@@ -184,8 +184,8 @@ class ConeType(enum.IntEnum):
     ELLIPTIC: elliptic
   """
 
-  PYRAMIDAL = mujoco.mjtCone.mjCONE_PYRAMIDAL
-  ELLIPTIC = mujoco.mjtCone.mjCONE_ELLIPTIC
+  PYRAMIDAL = mujoco.mjtCone.mjCONE_PYRAMIDAL  # pyrefly: ignore[bad-assignment]
+  ELLIPTIC = mujoco.mjtCone.mjCONE_ELLIPTIC  # pyrefly: ignore[bad-assignment]
 
 
 class JacobianType(enum.IntEnum):
@@ -197,9 +197,9 @@ class JacobianType(enum.IntEnum):
     AUTO: sparse if nv>60 and device is TPU, dense otherwise
   """
 
-  DENSE = mujoco.mjtJacobian.mjJAC_DENSE
-  SPARSE = mujoco.mjtJacobian.mjJAC_SPARSE
-  AUTO = mujoco.mjtJacobian.mjJAC_AUTO
+  DENSE = mujoco.mjtJacobian.mjJAC_DENSE  # pyrefly: ignore[bad-assignment]
+  SPARSE = mujoco.mjtJacobian.mjJAC_SPARSE  # pyrefly: ignore[bad-assignment]
+  AUTO = mujoco.mjtJacobian.mjJAC_AUTO  # pyrefly: ignore[bad-assignment]
 
 
 class SolverType(enum.IntEnum):
@@ -211,8 +211,8 @@ class SolverType(enum.IntEnum):
   """
 
   # unsupported: PGS
-  CG = mujoco.mjtSolver.mjSOL_CG
-  NEWTON = mujoco.mjtSolver.mjSOL_NEWTON
+  CG = mujoco.mjtSolver.mjSOL_CG  # pyrefly: ignore[bad-assignment]
+  NEWTON = mujoco.mjtSolver.mjSOL_NEWTON  # pyrefly: ignore[bad-assignment]
 
 
 class EqType(enum.IntEnum):
@@ -225,10 +225,10 @@ class EqType(enum.IntEnum):
     TENDON: couple the lengths of two tendons with cubic
   """
 
-  CONNECT = mujoco.mjtEq.mjEQ_CONNECT
-  WELD = mujoco.mjtEq.mjEQ_WELD
-  JOINT = mujoco.mjtEq.mjEQ_JOINT
-  TENDON = mujoco.mjtEq.mjEQ_TENDON
+  CONNECT = mujoco.mjtEq.mjEQ_CONNECT  # pyrefly: ignore[bad-assignment]
+  WELD = mujoco.mjtEq.mjEQ_WELD  # pyrefly: ignore[bad-assignment]
+  JOINT = mujoco.mjtEq.mjEQ_JOINT  # pyrefly: ignore[bad-assignment]
+  TENDON = mujoco.mjtEq.mjEQ_TENDON  # pyrefly: ignore[bad-assignment]
   # unsupported: DISTANCE
 
 
@@ -243,11 +243,11 @@ class WrapType(enum.IntEnum):
     CYLINDER: wrap around (infinite) cylinder
   """
 
-  JOINT = mujoco.mjtWrap.mjWRAP_JOINT
-  PULLEY = mujoco.mjtWrap.mjWRAP_PULLEY
-  SITE = mujoco.mjtWrap.mjWRAP_SITE
-  SPHERE = mujoco.mjtWrap.mjWRAP_SPHERE
-  CYLINDER = mujoco.mjtWrap.mjWRAP_CYLINDER
+  JOINT = mujoco.mjtWrap.mjWRAP_JOINT  # pyrefly: ignore[bad-assignment]
+  PULLEY = mujoco.mjtWrap.mjWRAP_PULLEY  # pyrefly: ignore[bad-assignment]
+  SITE = mujoco.mjtWrap.mjWRAP_SITE  # pyrefly: ignore[bad-assignment]
+  SPHERE = mujoco.mjtWrap.mjWRAP_SPHERE  # pyrefly: ignore[bad-assignment]
+  CYLINDER = mujoco.mjtWrap.mjWRAP_CYLINDER  # pyrefly: ignore[bad-assignment]
 
 
 class TrnType(enum.IntEnum):
@@ -260,10 +260,10 @@ class TrnType(enum.IntEnum):
     SITE: force on site
   """
 
-  JOINT = mujoco.mjtTrn.mjTRN_JOINT
-  JOINTINPARENT = mujoco.mjtTrn.mjTRN_JOINTINPARENT
-  SITE = mujoco.mjtTrn.mjTRN_SITE
-  TENDON = mujoco.mjtTrn.mjTRN_TENDON
+  JOINT = mujoco.mjtTrn.mjTRN_JOINT  # pyrefly: ignore[bad-assignment]
+  JOINTINPARENT = mujoco.mjtTrn.mjTRN_JOINTINPARENT  # pyrefly: ignore[bad-assignment]
+  SITE = mujoco.mjtTrn.mjTRN_SITE  # pyrefly: ignore[bad-assignment]
+  TENDON = mujoco.mjtTrn.mjTRN_TENDON  # pyrefly: ignore[bad-assignment]
   # unsupported: SLIDERCRANK, BODY
 
 
@@ -278,11 +278,11 @@ class DynType(enum.IntEnum):
     MUSCLE: piece-wise linear filter with two time constants
   """
 
-  NONE = mujoco.mjtDyn.mjDYN_NONE
-  INTEGRATOR = mujoco.mjtDyn.mjDYN_INTEGRATOR
-  FILTER = mujoco.mjtDyn.mjDYN_FILTER
-  FILTEREXACT = mujoco.mjtDyn.mjDYN_FILTEREXACT
-  MUSCLE = mujoco.mjtDyn.mjDYN_MUSCLE
+  NONE = mujoco.mjtDyn.mjDYN_NONE  # pyrefly: ignore[bad-assignment]
+  INTEGRATOR = mujoco.mjtDyn.mjDYN_INTEGRATOR  # pyrefly: ignore[bad-assignment]
+  FILTER = mujoco.mjtDyn.mjDYN_FILTER  # pyrefly: ignore[bad-assignment]
+  FILTEREXACT = mujoco.mjtDyn.mjDYN_FILTEREXACT  # pyrefly: ignore[bad-assignment]
+  MUSCLE = mujoco.mjtDyn.mjDYN_MUSCLE  # pyrefly: ignore[bad-assignment]
   # unsupported: USER
 
 
@@ -295,9 +295,9 @@ class GainType(enum.IntEnum):
     MUSCLE: muscle FLV curve computed by muscle_gain
   """
 
-  FIXED = mujoco.mjtGain.mjGAIN_FIXED
-  AFFINE = mujoco.mjtGain.mjGAIN_AFFINE
-  MUSCLE = mujoco.mjtGain.mjGAIN_MUSCLE
+  FIXED = mujoco.mjtGain.mjGAIN_FIXED  # pyrefly: ignore[bad-assignment]
+  AFFINE = mujoco.mjtGain.mjGAIN_AFFINE  # pyrefly: ignore[bad-assignment]
+  MUSCLE = mujoco.mjtGain.mjGAIN_MUSCLE  # pyrefly: ignore[bad-assignment]
   # unsupported: USER
 
 
@@ -310,9 +310,9 @@ class BiasType(enum.IntEnum):
     MUSCLE: muscle passive force computed by muscle_bias
   """
 
-  NONE = mujoco.mjtBias.mjBIAS_NONE
-  AFFINE = mujoco.mjtBias.mjBIAS_AFFINE
-  MUSCLE = mujoco.mjtBias.mjBIAS_MUSCLE
+  NONE = mujoco.mjtBias.mjBIAS_NONE  # pyrefly: ignore[bad-assignment]
+  AFFINE = mujoco.mjtBias.mjBIAS_AFFINE  # pyrefly: ignore[bad-assignment]
+  MUSCLE = mujoco.mjtBias.mjBIAS_MUSCLE  # pyrefly: ignore[bad-assignment]
   # unsupported: USER
 
 
@@ -327,14 +327,14 @@ class ConstraintType(enum.IntEnum):
     CONTACT_PYRAMIDAL: frictional contact, pyramidal friction cone
   """
 
-  EQUALITY = mujoco.mjtConstraint.mjCNSTR_EQUALITY
-  FRICTION_DOF = mujoco.mjtConstraint.mjCNSTR_FRICTION_DOF
-  FRICTION_TENDON = mujoco.mjtConstraint.mjCNSTR_FRICTION_TENDON
-  LIMIT_JOINT = mujoco.mjtConstraint.mjCNSTR_LIMIT_JOINT
-  LIMIT_TENDON = mujoco.mjtConstraint.mjCNSTR_LIMIT_TENDON
-  CONTACT_FRICTIONLESS = mujoco.mjtConstraint.mjCNSTR_CONTACT_FRICTIONLESS
-  CONTACT_PYRAMIDAL = mujoco.mjtConstraint.mjCNSTR_CONTACT_PYRAMIDAL
-  CONTACT_ELLIPTIC = mujoco.mjtConstraint.mjCNSTR_CONTACT_ELLIPTIC
+  EQUALITY = mujoco.mjtConstraint.mjCNSTR_EQUALITY  # pyrefly: ignore[bad-assignment]
+  FRICTION_DOF = mujoco.mjtConstraint.mjCNSTR_FRICTION_DOF  # pyrefly: ignore[bad-assignment]
+  FRICTION_TENDON = mujoco.mjtConstraint.mjCNSTR_FRICTION_TENDON  # pyrefly: ignore[bad-assignment]
+  LIMIT_JOINT = mujoco.mjtConstraint.mjCNSTR_LIMIT_JOINT  # pyrefly: ignore[bad-assignment]
+  LIMIT_TENDON = mujoco.mjtConstraint.mjCNSTR_LIMIT_TENDON  # pyrefly: ignore[bad-assignment]
+  CONTACT_FRICTIONLESS = mujoco.mjtConstraint.mjCNSTR_CONTACT_FRICTIONLESS  # pyrefly: ignore[bad-assignment]
+  CONTACT_PYRAMIDAL = mujoco.mjtConstraint.mjCNSTR_CONTACT_PYRAMIDAL  # pyrefly: ignore[bad-assignment]
+  CONTACT_ELLIPTIC = mujoco.mjtConstraint.mjCNSTR_CONTACT_ELLIPTIC  # pyrefly: ignore[bad-assignment]
 
 
 class CamLightType(enum.IntEnum):
@@ -348,11 +348,11 @@ class CamLightType(enum.IntEnum):
     TARGETBODYCOM: pos fixed in body, rot tracks target subtree com
   """
 
-  FIXED = mujoco.mjtCamLight.mjCAMLIGHT_FIXED
-  TRACK = mujoco.mjtCamLight.mjCAMLIGHT_TRACK
-  TRACKCOM = mujoco.mjtCamLight.mjCAMLIGHT_TRACKCOM
-  TARGETBODY = mujoco.mjtCamLight.mjCAMLIGHT_TARGETBODY
-  TARGETBODYCOM = mujoco.mjtCamLight.mjCAMLIGHT_TARGETBODYCOM
+  FIXED = mujoco.mjtCamLight.mjCAMLIGHT_FIXED  # pyrefly: ignore[bad-assignment]
+  TRACK = mujoco.mjtCamLight.mjCAMLIGHT_TRACK  # pyrefly: ignore[bad-assignment]
+  TRACKCOM = mujoco.mjtCamLight.mjCAMLIGHT_TRACKCOM  # pyrefly: ignore[bad-assignment]
+  TARGETBODY = mujoco.mjtCamLight.mjCAMLIGHT_TARGETBODY  # pyrefly: ignore[bad-assignment]
+  TARGETBODYCOM = mujoco.mjtCamLight.mjCAMLIGHT_TARGETBODYCOM  # pyrefly: ignore[bad-assignment]
 
 
 class SensorType(enum.IntEnum):
diff --git a/mjx/mujoco/mjx/codegen/generate_warp_shim.py b/mjx/mujoco/mjx/codegen/generate_warp_shim.py
index a513b88c..46516f88 100644
--- a/mjx/mujoco/mjx/codegen/generate_warp_shim.py
+++ b/mjx/mujoco/mjx/codegen/generate_warp_shim.py
@@ -264,8 +264,8 @@ def _warp_function(
   # create a dummy output if there are no output fields
   needs_dummy_output = not field_usage.data_out_fields
   if needs_dummy_output and fn_name != 'render':
-    fn_args_raw.append('# Dummy output')
-    fn_args_raw.append('dummy: wp.array[int],')
+    fn_args_raw.append('# Dummy output')  # pyrefly: ignore[bad-argument-type]
+    fn_args_raw.append('dummy: wp.array[int],')  # pyrefly: ignore[bad-argument-type]
 
   return fn_args_raw, fn_assignments, fn_call
 
diff --git a/mjx/mujoco/mjx/codegen/generate_warp_types.py b/mjx/mujoco/mjx/codegen/generate_warp_types.py
index 48e090e7..9f930611 100644
--- a/mjx/mujoco/mjx/codegen/generate_warp_types.py
+++ b/mjx/mujoco/mjx/codegen/generate_warp_types.py
@@ -440,7 +440,7 @@ def write_core_cls(
     shape_property = _DATA_SHAPE_PROPERTY_FIELD
 
   new_class_body = _build_new_class_body_ast(
-      keys,
+      keys,  # pyrefly: ignore[bad-argument-type]
       cls_name,
       annotations,
       shape_property=shape_property,
diff --git a/mjx/mujoco/mjx/codegen/trace.py b/mjx/mujoco/mjx/codegen/trace.py
index b62896a3..f8339aa1 100644
--- a/mjx/mujoco/mjx/codegen/trace.py
+++ b/mjx/mujoco/mjx/codegen/trace.py
@@ -162,7 +162,7 @@ class _FunctionFieldUsageVisitor(ast.NodeVisitor):
       called_fn_name = node.func.id
       key = (hash(self._current_fpath), called_fn_name)
       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(
             called_fn_name, self._current_fpath
         )
@@ -181,8 +181,8 @@ class _FunctionFieldUsageVisitor(ast.NodeVisitor):
         if len(node.args) != 2:
           raise ValueError(f'wp.copy() must have 2 arguments, got {node.args}.')
         out_node, in_node = node.args
-        self.add_field_usage(out_node, True)
-        self.add_field_usage(in_node, False)
+        self.add_field_usage(out_node, True)  # pyrefly: ignore[bad-argument-type]
+        self.add_field_usage(in_node, False)  # pyrefly: ignore[bad-argument-type]
         return
 
       # do not trace array creation
@@ -201,7 +201,7 @@ class _FunctionFieldUsageVisitor(ast.NodeVisitor):
       key = (hash(self._current_fpath), called_fn_name)
       next_fpath = self._module_fpaths.get(parts[0])
       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.generic_visit(node)
@@ -223,10 +223,10 @@ def trace_function(
 ) -> FieldUsage:
   """Traces the function statically to find usages of model and data fields."""
   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)
 
-  src = fpath.read_text()
+  src = fpath.read_text()  # pyrefly: ignore[missing-attribute]
   parsed_ast = ast.parse(src, filename=str(fpath))
 
   target_fn_nodes = (
@@ -262,7 +262,7 @@ def trace_function(
   if visited_fns is None:
     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:
     visitor.visit(body)
 
diff --git a/mjx/mujoco/mjx/warp/ffi.py b/mjx/mujoco/mjx/warp/ffi.py
index be7b1468..b6e0998a 100644
--- a/mjx/mujoco/mjx/warp/ffi.py
+++ b/mjx/mujoco/mjx/warp/ffi.py
@@ -119,7 +119,7 @@ def jax_callable_variadic_tuple(
 
     # Provide a flattened signature for the Warp callable machinery.
     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__ = {
         p.name: p.annotation
         for p in new_signature.parameters.values()
@@ -205,7 +205,7 @@ def _expand_dim_from_path(
   if ndim is None or ndim < 0:
     return leaf
   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:
     raise AssertionError(
         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}'
     )
   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
 
 
@@ -288,7 +288,7 @@ def _maybe_broadcast_to(
   """Broadcasts fields that are used in MuJoCo Warp."""
   ndim = _get_mapping_from_tree_path(path, mjx_warp_types._NDIM[cls_str])
   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)
   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."""
   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)
   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(
         # NB: if a field is not in MuJoCo Warp, we let JAX do its magic.
         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,
         out_batched,
diff --git a/python/mujoco/__init__.py b/python/mujoco/__init__.py
index fdd86372..1147be58 100644
--- a/python/mujoco/__init__.py
+++ b/python/mujoco/__init__.py
@@ -232,8 +232,8 @@ def _bind_data(
 
 _specs.MjSpec.from_zip = from_zip
 _specs.MjSpec.to_zip = to_zip
-_structs.MjData.bind = _bind_data
-_structs.MjModel.bind = _bind_model
+_structs.MjData.bind = _bind_data  # pyrefly: ignore[bad-assignment]
+_structs.MjModel.bind = _bind_model  # pyrefly: ignore[bad-assignment]
 
 HEADERS_DIR = os.path.join(os.path.dirname(__file__), 'include/mujoco')
 PLUGINS_DIR = os.path.join(os.path.dirname(__file__), 'plugin')
diff --git a/python/mujoco/codegen/generate_spec_bindings.py b/python/mujoco/codegen/generate_spec_bindings.py
index 2cef63d5..c3d23423 100644
--- a/python/mujoco/codegen/generate_spec_bindings.py
+++ b/python/mujoco/codegen/generate_spec_bindings.py
@@ -142,10 +142,10 @@ def _struct_binding_code(
       for f 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
-  field = ast_nodes.ValueType(name=name)
-  code += _value_binding_code(field, classname, varname)
+  field = ast_nodes.ValueType(name=name)  # pyrefly: ignore[bad-assignment]
+  code += _value_binding_code(field, classname, varname)  # pyrefly: ignore[bad-argument-type]
   return code
 
 
@@ -211,7 +211,7 @@ def _ptr_binding_code(
   if vartype == 'mjsElement':  # this is ignored by the caller
     return 'mjsElement'
   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
     return f"""\
   {classname}.def_property(
@@ -347,7 +347,7 @@ def generate() -> None:
     ) and key != 'mjsElement':
       print('\n  // ' + key)
       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':
           print(code)
 
@@ -556,7 +556,7 @@ def generate_add() -> None:
       py_args = ['py::arg("name") = py::none()']
 
     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:
         code_field = code_field + '\n        ' + line
         set_types.append(set_type)
diff --git a/python/mujoco/introspect/codegen/generate.py b/python/mujoco/introspect/codegen/generate.py
index 06e07a0c..1a48eea6 100644
--- a/python/mujoco/introspect/codegen/generate.py
+++ b/python/mujoco/introspect/codegen/generate.py
@@ -126,22 +126,22 @@ def main(argv: Sequence[str]) -> None:
     raise app.UsageError('Too many command-line arguments.')
 
   visitor = ast_processor.process(
-      _JSON_PATH.value,
+      _JSON_PATH.value,  # pyrefly: ignore[bad-argument-type]
       (_HEADER_PATHS.value or '').split(),
       _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)
     f.write(_HEADER_TEMPLATE.format(year=2022, type='functions'))
     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)
     f.write(_HEADER_TEMPLATE.format(year=2022, type='enums'))
     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)
     f.write(_HEADER_TEMPLATE.format(year=2023, type='structs'))
     f.write(_STRUCTS_TEMPLATE.format(structs_str=structs_str))
diff --git a/python/mujoco/introspect/type_parsing.py b/python/mujoco/introspect/type_parsing.py
index dfc2ab53..8cb2174b 100644
--- a/python/mujoco/introspect/type_parsing.py
+++ b/python/mujoco/introspect/type_parsing.py
@@ -57,7 +57,7 @@ def _parse_maybe_array(
         int(s.strip()) for s in ARRAY_N_PATTERN.findall(array_match.group(0)))
     inner_type_str = type_name[:array_match.start()]
     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)
   else:
     return _parse_maybe_pointer(type_name, innermost_type)
@@ -140,7 +140,7 @@ def parse_type(
   result = None
   while type_str_stack:
     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:
       raise ValueError(f'invalid type name {type_name!r}') from e
 
diff --git a/python/mujoco/minimize.py b/python/mujoco/minimize.py
index 5064efb0..955744d6 100644
--- a/python/mujoco/minimize.py
+++ b/python/mujoco/minimize.py
@@ -405,7 +405,7 @@ def least_squares(
       print(message, file=output)
 
     # 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:
       log = dataclasses.replace(
           log, residual=r, jacobian=jac, grad=grad / D, step=step
diff --git a/python/mujoco/rendering/classic/renderer.py b/python/mujoco/rendering/classic/renderer.py
index 0e472e58..d91c7c46 100644
--- a/python/mujoco/rendering/classic/renderer.py
+++ b/python/mujoco/rendering/classic/renderer.py
@@ -93,7 +93,7 @@ the clause:
     mujoco.mjr_setBuffer(
         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.
     self._depth_rendering = False
@@ -291,10 +291,10 @@ the clause:
       # Defaults to mjCAMERA_FREE, otherwise mjCAMERA_FIXED refers to a
       # camera explicitly defined in the model.
       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)
       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
     mujoco.mjv_updateScene(
diff --git a/python/mujoco/rollout.py b/python/mujoco/rollout.py
index daf55504..5e0ebf80 100644
--- a/python/mujoco/rollout.py
+++ b/python/mujoco/rollout.py
@@ -180,7 +180,7 @@ class Rollout:
       model = [model]  # Use a length 1 list to simplify code below
 
     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
     nstep = _infer_dimension(
@@ -310,9 +310,9 @@ def rollout(
     ValueError: bad shapes or sizes.
   """  # fmt: skip
   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
   if persistent_pool:
diff --git a/python/mujoco/sysid/_src/optimize.py b/python/mujoco/sysid/_src/optimize.py
index b0e22982..7a0209af 100644
--- a/python/mujoco/sysid/_src/optimize.py
+++ b/python/mujoco/sysid/_src/optimize.py
@@ -53,7 +53,7 @@ def _warn_if_ill_conditioned(
           residual=f,
           x=x0.reshape(-1, 1),
           r=r0,
-          eps=eps,
+          eps=eps,  # pyrefly: ignore[bad-argument-type]
           n_res=0,
           bounds=[bounds[0].reshape(-1, 1), bounds[1].reshape(-1, 1)],
       )[0],
diff --git a/python/mujoco/viewer.py b/python/mujoco/viewer.py
index 7140773d..e75d295d 100644
--- a/python/mujoco/viewer.py
+++ b/python/mujoco/viewer.py
@@ -325,8 +325,8 @@ def _reload(
 
 def _physics_loop(simulate: _Simulate, loader: Optional[_InternalLoaderType]):
   """Physics loop for the GUI, to be run in a separate thread."""
-  m: mujoco.MjModel = None
-  d: mujoco.MjData = None
+  m: mujoco.MjModel = None  # pyrefly: ignore[bad-assignment]
+  d: mujoco.MjData = None  # pyrefly: ignore[bad-assignment]
   ctrl_noise = np.array([])
   reload = True