diff --git a/mjx/mujoco/mjx/_src/render.py b/mjx/mujoco/mjx/_src/render.py index 4ba7e215..bc5fc044 100644 --- a/mjx/mujoco/mjx/_src/render.py +++ b/mjx/mujoco/mjx/_src/render.py @@ -53,7 +53,7 @@ def _call_render( token = token_array.reshape( d._impl._jax_token.shape # pytype: disable=attribute-error ) - d = d.tree_replace({'_impl._jax_token': token}) + d = d.tree_replace({'_impl._jax_token': token}) # pyrefly: ignore[bad-assignment] return rgb, depth, seg, d raise NotImplementedError('render only implemented for MuJoCo Warp.') diff --git a/mjx/mujoco/mjx/_src/sensor.py b/mjx/mujoco/mjx/_src/sensor.py index d3ae2da8..01587be1 100644 --- a/mjx/mujoco/mjx/_src/sensor.py +++ b/mjx/mujoco/mjx/_src/sensor.py @@ -466,7 +466,7 @@ def sensor_acc(m: Model, d: Data) -> Data: # compute contact forces contact_force = [] condim_ids = [] - for dim in set(d._impl.contact.dim): + for dim in set(d._impl.contact.dim): # pyrefly: ignore[missing-attribute] force, condim_id = support.contact_force_dim(m, d, dim) contact_force.append(force) condim_ids.append(condim_id) @@ -485,7 +485,7 @@ def sensor_acc(m: Model, d: Data) -> Data: if sensor_type == SensorType.TOUCH: # get bodies of contact geoms - conbody = jp.array(m.geom_bodyid)[d._impl.contact.geom] + conbody = jp.array(m.geom_bodyid)[d._impl.contact.geom] # pyrefly: ignore[missing-attribute] # get site information site_bodyid = m.site_bodyid[objid] @@ -495,14 +495,14 @@ def sensor_acc(m: Model, d: Data) -> Data: site_type = m.site_type[objid] conbody0 = site_bodyid[:, None] == conbody[:, 0] conbody1 = site_bodyid[:, None] == conbody[:, 1] - contacts = (d._impl.contact.efc_address >= 0)[None] & ( + contacts = (d._impl.contact.efc_address >= 0)[None] & ( # pyrefly: ignore[missing-attribute] conbody0 | conbody1 ) # compute conray, flip if second body conray = jax.vmap( lambda frame, force: math.normalize(frame[0] * force[0]) - )(d._impl.contact.frame, contact_force) # pyrefly: ignore[unbound-name] + )(d._impl.contact.frame, contact_force) # pyrefly: ignore[missing-attribute, unbound-name] conray = jp.where(conbody1[..., None], -conray, conray) # compute distance, mapping over sites and contacts @@ -524,7 +524,7 @@ def sensor_acc(m: Model, d: Data) -> Data: site_xpos[dist_id_site], site_xmat[dist_id_site], st, - d._impl.contact.pos, + d._impl.contact.pos, # pyrefly: ignore[missing-attribute] conray[dist_id_site], ) dist.append(jp.where(jp.isinf(dist_site), 0, dist_site)) @@ -535,11 +535,11 @@ def sensor_acc(m: Model, d: Data) -> Data: sensor = jp.dot((dist > 0) & contacts, contact_force[:, 0]) elif sensor_type == SensorType.CONTACT: # maximum number of contacts - ncon = d._impl.ncon + ncon = d._impl.ncon # pyrefly: ignore[missing-attribute] # active contacts - dist = d._impl.contact.dist - pos = dist - d._impl.contact.includemargin + dist = d._impl.contact.dist # pyrefly: ignore[missing-attribute] + pos = dist - d._impl.contact.includemargin # pyrefly: ignore[missing-attribute] is_contact = pos < 0 # reduction criteria @@ -607,8 +607,8 @@ def sensor_acc(m: Model, d: Data) -> Data: elif objtype == ObjType.GEOM or reftype == ObjType.GEOM: sensorid1 = objid[idx_ds] sensorid2 = refid[idx_ds] - geomid0 = d._impl.contact.geom[:, 0] - geomid1 = d._impl.contact.geom[:, 1] + geomid0 = d._impl.contact.geom[:, 0] # pyrefly: ignore[missing-attribute] + geomid1 = d._impl.contact.geom[:, 1] # pyrefly: ignore[missing-attribute] # match sensor ids and contact geom ids geom0id1 = geomid0 == sensorid1[:, None] @@ -671,13 +671,13 @@ def sensor_acc(m: Model, d: Data) -> Data: slot.append(dist[cid, None]) if dataspec & (1 << 4): # pos - slot.append(d._impl.contact.pos[cid]) + slot.append(d._impl.contact.pos[cid]) # pyrefly: ignore[missing-attribute] if dataspec & (1 << 5): # normal - slot.append(flip[:, 2, None] * d._impl.contact.frame[cid, 0]) + slot.append(flip[:, 2, None] * d._impl.contact.frame[cid, 0]) # pyrefly: ignore[missing-attribute] if dataspec & (1 << 6): # tangent - slot.append(flip[:, 2, None] * d._impl.contact.frame[cid, 1]) + slot.append(flip[:, 2, None] * d._impl.contact.frame[cid, 1]) # pyrefly: ignore[missing-attribute] found = jp.tile(jp.arange(num), nsensor) < jp.repeat(nfound, num) sensors.append((found[:, None] * jp.hstack(slot)).reshape(-1)) diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index 50631dbe..ad4d25ee 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/mjx/mujoco/mjx/_src/types.py @@ -395,40 +395,40 @@ class SensorType(enum.IntEnum): FRAMEANGACC: 3D angular acceleration """ - MAGNETOMETER = mujoco.mjtSensor.mjSENS_MAGNETOMETER - CAMPROJECTION = mujoco.mjtSensor.mjSENS_CAMPROJECTION - RANGEFINDER = mujoco.mjtSensor.mjSENS_RANGEFINDER - JOINTPOS = mujoco.mjtSensor.mjSENS_JOINTPOS - TENDONPOS = mujoco.mjtSensor.mjSENS_TENDONPOS - ACTUATORPOS = mujoco.mjtSensor.mjSENS_ACTUATORPOS - BALLQUAT = mujoco.mjtSensor.mjSENS_BALLQUAT - FRAMEPOS = mujoco.mjtSensor.mjSENS_FRAMEPOS - FRAMEXAXIS = mujoco.mjtSensor.mjSENS_FRAMEXAXIS - FRAMEYAXIS = mujoco.mjtSensor.mjSENS_FRAMEYAXIS - FRAMEZAXIS = mujoco.mjtSensor.mjSENS_FRAMEZAXIS - FRAMEQUAT = mujoco.mjtSensor.mjSENS_FRAMEQUAT - SUBTREECOM = mujoco.mjtSensor.mjSENS_SUBTREECOM - CLOCK = mujoco.mjtSensor.mjSENS_CLOCK - VELOCIMETER = mujoco.mjtSensor.mjSENS_VELOCIMETER - GYRO = mujoco.mjtSensor.mjSENS_GYRO - JOINTVEL = mujoco.mjtSensor.mjSENS_JOINTVEL - TENDONVEL = mujoco.mjtSensor.mjSENS_TENDONVEL - ACTUATORVEL = mujoco.mjtSensor.mjSENS_ACTUATORVEL - BALLANGVEL = mujoco.mjtSensor.mjSENS_BALLANGVEL - FRAMELINVEL = mujoco.mjtSensor.mjSENS_FRAMELINVEL - FRAMEANGVEL = mujoco.mjtSensor.mjSENS_FRAMEANGVEL - SUBTREELINVEL = mujoco.mjtSensor.mjSENS_SUBTREELINVEL - SUBTREEANGMOM = mujoco.mjtSensor.mjSENS_SUBTREEANGMOM - TOUCH = mujoco.mjtSensor.mjSENS_TOUCH - CONTACT = mujoco.mjtSensor.mjSENS_CONTACT - ACCELEROMETER = mujoco.mjtSensor.mjSENS_ACCELEROMETER - FORCE = mujoco.mjtSensor.mjSENS_FORCE - TORQUE = mujoco.mjtSensor.mjSENS_TORQUE - ACTUATORFRC = mujoco.mjtSensor.mjSENS_ACTUATORFRC - JOINTACTFRC = mujoco.mjtSensor.mjSENS_JOINTACTFRC - TENDONACTFRC = mujoco.mjtSensor.mjSENS_TENDONACTFRC - FRAMELINACC = mujoco.mjtSensor.mjSENS_FRAMELINACC - FRAMEANGACC = mujoco.mjtSensor.mjSENS_FRAMEANGACC + MAGNETOMETER = mujoco.mjtSensor.mjSENS_MAGNETOMETER # pyrefly: ignore[bad-assignment] + CAMPROJECTION = mujoco.mjtSensor.mjSENS_CAMPROJECTION # pyrefly: ignore[bad-assignment] + RANGEFINDER = mujoco.mjtSensor.mjSENS_RANGEFINDER # pyrefly: ignore[bad-assignment] + JOINTPOS = mujoco.mjtSensor.mjSENS_JOINTPOS # pyrefly: ignore[bad-assignment] + TENDONPOS = mujoco.mjtSensor.mjSENS_TENDONPOS # pyrefly: ignore[bad-assignment] + ACTUATORPOS = mujoco.mjtSensor.mjSENS_ACTUATORPOS # pyrefly: ignore[bad-assignment] + BALLQUAT = mujoco.mjtSensor.mjSENS_BALLQUAT # pyrefly: ignore[bad-assignment] + FRAMEPOS = mujoco.mjtSensor.mjSENS_FRAMEPOS # pyrefly: ignore[bad-assignment] + FRAMEXAXIS = mujoco.mjtSensor.mjSENS_FRAMEXAXIS # pyrefly: ignore[bad-assignment] + FRAMEYAXIS = mujoco.mjtSensor.mjSENS_FRAMEYAXIS # pyrefly: ignore[bad-assignment] + FRAMEZAXIS = mujoco.mjtSensor.mjSENS_FRAMEZAXIS # pyrefly: ignore[bad-assignment] + FRAMEQUAT = mujoco.mjtSensor.mjSENS_FRAMEQUAT # pyrefly: ignore[bad-assignment] + SUBTREECOM = mujoco.mjtSensor.mjSENS_SUBTREECOM # pyrefly: ignore[bad-assignment] + CLOCK = mujoco.mjtSensor.mjSENS_CLOCK # pyrefly: ignore[bad-assignment] + VELOCIMETER = mujoco.mjtSensor.mjSENS_VELOCIMETER # pyrefly: ignore[bad-assignment] + GYRO = mujoco.mjtSensor.mjSENS_GYRO # pyrefly: ignore[bad-assignment] + JOINTVEL = mujoco.mjtSensor.mjSENS_JOINTVEL # pyrefly: ignore[bad-assignment] + TENDONVEL = mujoco.mjtSensor.mjSENS_TENDONVEL # pyrefly: ignore[bad-assignment] + ACTUATORVEL = mujoco.mjtSensor.mjSENS_ACTUATORVEL # pyrefly: ignore[bad-assignment] + BALLANGVEL = mujoco.mjtSensor.mjSENS_BALLANGVEL # pyrefly: ignore[bad-assignment] + FRAMELINVEL = mujoco.mjtSensor.mjSENS_FRAMELINVEL # pyrefly: ignore[bad-assignment] + FRAMEANGVEL = mujoco.mjtSensor.mjSENS_FRAMEANGVEL # pyrefly: ignore[bad-assignment] + SUBTREELINVEL = mujoco.mjtSensor.mjSENS_SUBTREELINVEL # pyrefly: ignore[bad-assignment] + SUBTREEANGMOM = mujoco.mjtSensor.mjSENS_SUBTREEANGMOM # pyrefly: ignore[bad-assignment] + TOUCH = mujoco.mjtSensor.mjSENS_TOUCH # pyrefly: ignore[bad-assignment] + CONTACT = mujoco.mjtSensor.mjSENS_CONTACT # pyrefly: ignore[bad-assignment] + ACCELEROMETER = mujoco.mjtSensor.mjSENS_ACCELEROMETER # pyrefly: ignore[bad-assignment] + FORCE = mujoco.mjtSensor.mjSENS_FORCE # pyrefly: ignore[bad-assignment] + TORQUE = mujoco.mjtSensor.mjSENS_TORQUE # pyrefly: ignore[bad-assignment] + ACTUATORFRC = mujoco.mjtSensor.mjSENS_ACTUATORFRC # pyrefly: ignore[bad-assignment] + JOINTACTFRC = mujoco.mjtSensor.mjSENS_JOINTACTFRC # pyrefly: ignore[bad-assignment] + TENDONACTFRC = mujoco.mjtSensor.mjSENS_TENDONACTFRC # pyrefly: ignore[bad-assignment] + FRAMELINACC = mujoco.mjtSensor.mjSENS_FRAMELINACC # pyrefly: ignore[bad-assignment] + FRAMEANGACC = mujoco.mjtSensor.mjSENS_FRAMEANGACC # pyrefly: ignore[bad-assignment] class ObjType(PyTreeNode): @@ -1213,6 +1213,6 @@ def tree_path_to_attr_str(path: jax.tree_util.KeyPath) -> str: path = path[: is_seq_key.index(True)] assert all(isinstance(p, jax.tree_util.GetAttrKey) for p in path) - path = [p for p in path if p.name != '_impl'] + path = [p for p in path if p.name != '_impl'] # pyrefly: ignore[bad-assignment] return '__'.join(p.name for p in path) diff --git a/mjx/mujoco/mjx/codegen/generate_warp_types.py b/mjx/mujoco/mjx/codegen/generate_warp_types.py index 1e3eebd1..09eaf088 100644 --- a/mjx/mujoco/mjx/codegen/generate_warp_types.py +++ b/mjx/mujoco/mjx/codegen/generate_warp_types.py @@ -134,7 +134,7 @@ def _get_target_annotation_node( return _ast_parse_type('int') if dataclasses.is_dataclass(annotation): - return _ast_parse_type(annotation.__name__) + return _ast_parse_type(annotation.__name__) # pyrefly: ignore[missing-attribute] is_tuple = typing.get_origin(annotation) == tuple if is_tuple and typing.get_args(annotation)[1] != ...: @@ -230,7 +230,7 @@ def _build_new_class_body_ast( ) new_body_nodes.append( - ast.AnnAssign( + ast.AnnAssign( # pyrefly: ignore[no-matching-overload] target=ast.Name(id=key, ctx=ast.Store()), annotation=annotation_node, value=value_node, diff --git a/mjx/mujoco/mjx/codegen/trace.py b/mjx/mujoco/mjx/codegen/trace.py index 0207652f..7fad9621 100644 --- a/mjx/mujoco/mjx/codegen/trace.py +++ b/mjx/mujoco/mjx/codegen/trace.py @@ -69,7 +69,7 @@ def _get_imported_module_fpaths(fpath: epath.Path) -> Dict[str, str]: if not (fully_qualified_name.startswith('mujoco_warp') or fully_qualified_name.startswith('mujoco')): continue - fpath = _resolve_module_name_to_fpath(fully_qualified_name) + fpath = _resolve_module_name_to_fpath(fully_qualified_name) # pyrefly: ignore[bad-assignment] if fpath: all_resolved_fpaths[alias] = fpath return all_resolved_fpaths diff --git a/python/mujoco/experimental/studio/sample/ghost.py b/python/mujoco/experimental/studio/sample/ghost.py index 213b6450..f2da0955 100644 --- a/python/mujoco/experimental/studio/sample/ghost.py +++ b/python/mujoco/experimental/studio/sample/ghost.py @@ -117,7 +117,7 @@ class GhostRenderer: self._viewer.extra_geoms.clear() if self._last_time is None or data.time < self._last_time: self._history.clear() - if not self._history or data.time > self._last_time: + if not self._history or data.time > self._last_time: # pyrefly: ignore[unsupported-operation] self._history.append(( data.time, data.geom_xpos.copy(), @@ -129,7 +129,7 @@ class GhostRenderer: while len(self._history) > 1 and self._history[1][0] <= target_time: self._history.popleft() - _, xpos, xmat = self._history[0] + _, xpos, xmat = self._history[0] # pyrefly: ignore[bad-assignment] if len(xpos) != model.ngeom or len(xmat) != model.ngeom: return diff --git a/python/mujoco/experimental/studio/sample/render.py b/python/mujoco/experimental/studio/sample/render.py index 07594b6c..f3fd6b81 100644 --- a/python/mujoco/experimental/studio/sample/render.py +++ b/python/mujoco/experimental/studio/sample/render.py @@ -70,7 +70,7 @@ def main(argv: list[str]) -> None: print(f'Error saving image to `{_OUTPUT.value}`: {ex}') sys.exit(-3) - return 0 + return 0 # pyrefly: ignore[bad-return] if __name__ == '__main__': diff --git a/python/mujoco/experimental/studio/viewer.py b/python/mujoco/experimental/studio/viewer.py index 4cda28f6..5ca3f365 100644 --- a/python/mujoco/experimental/studio/viewer.py +++ b/python/mujoco/experimental/studio/viewer.py @@ -59,7 +59,7 @@ def main(argv: list[str]) -> None: ) as handle: # Send the model to the viewer, if we have a model. if model is not None: - handle.send_to_viewer(messages.ModelEvent(model=model, path=model_path)) + handle.send_to_viewer(messages.ModelEvent(model=model, path=model_path)) # pyrefly: ignore[bad-argument-type] # Run the simulation. step_control = sim.StepControl() diff --git a/python/mujoco/experimental/studio/web/web_server.py b/python/mujoco/experimental/studio/web/web_server.py index 3266af87..b4dbab53 100644 --- a/python/mujoco/experimental/studio/web/web_server.py +++ b/python/mujoco/experimental/studio/web/web_server.py @@ -792,7 +792,7 @@ def _run_server( try: # Handshake: browser CmdVersion -> client; client CmdVersion -> browser. browser_version = await ws.recv() - my_writer.write(browser_version) + my_writer.write(browser_version) # pyrefly: ignore[bad-argument-type] await my_writer.drain() server_version = await my_reader.readexactly(_NETIMGUI_CMD_VERSION_SIZE) await ws.send(server_version) @@ -869,10 +869,10 @@ def _run_server( gen_before = generation.value if gen_before & 1: continue - (used,) = struct.unpack(" shm_capacity: return None - data = bytes(shm_array[4 : 4 + used]) + data = bytes(shm_array[4 : 4 + used]) # pyrefly: ignore[unsupported-operation] if generation.value == gen_before: return gen_before, data return None diff --git a/python/mujoco/usd/demo.py b/python/mujoco/usd/demo.py index d378df4e..94d669ba 100644 --- a/python/mujoco/usd/demo.py +++ b/python/mujoco/usd/demo.py @@ -31,7 +31,7 @@ def generate_usd_trajectory(local_args): # create an instance of the USDExporter exp = exporter.USDExporter( model=m, - output_directory_name=pathlib.Path(local_args.model_path).stem, + output_directory_name=pathlib.Path(local_args.model_path).stem, # pyrefly: ignore[unexpected-keyword] output_directory_root=local_args.output_directory_root, camera_names=local_args.camera_names, ) @@ -42,7 +42,7 @@ def generate_usd_trajectory(local_args): if exp.frame_count < d.time * local_args.framerate: exp.update_scene(data=d) - exp.add_light(pos=(0, 0, 0), + exp.add_light(pos=(0, 0, 0), # pyrefly: ignore[bad-argument-type] intensity=2000, light_type='dome')