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
|
||||
|
||||
Reference in New Issue
Block a user