suppress pyrefly errors in mujoco
Startblock: * // Put other blockers before this line to avoid churning. * has lgtm * is approved * and then * all comments are resolved * and then PiperOrigin-RevId: 961854282 Change-Id: I9d323a8ef9f36ffe69cc62d449849eac82dc6858
This commit is contained in:
Hana Joo
committed by
Copybara-Service
Copybara-Service
parent
d6dc966c04
commit
0aad65868d
@@ -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.')
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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__':
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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("<I", bytes(shm_array[:4]))
|
||||
(used,) = struct.unpack("<I", bytes(shm_array[:4])) # pyrefly: ignore[unsupported-operation]
|
||||
if used == 0 or used > 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
|
||||
|
||||
@@ -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')
|
||||
|
||||
|
||||
Reference in New Issue
Block a user