ci: fix stubgen output and mypy lint

Author: juelg

Makefile

@@ -33,6 +33,8 @@ stubgen:
 	find ./python/rcs/_core -name '*.pyi' -print | xargs sed -i 's/tuple\[typing\.Literal\[\([0-9]\+\)\], typing\.Literal\[1\]\]/tuple\[typing\.Literal[\1]\]/g'
 	find ./python/rcs/_core -name '*.pyi' -print | xargs sed -i 's/tuple\[\([M|N]\), typing\.Literal\[1\]\]/tuple\[\1\]/g'
 	sed -i 's/    q_home: numpy\.ndarray\[tuple\[M\], numpy\.dtype\[numpy\.float64\]\] | None/    q_home: numpy.ndarray | None/' python/rcs/_core/common.pyi
+	python -c "from pathlib import Path; p=Path('python/rcs/_core/common.pyi'); lines=p.read_text().splitlines(); lines=[line for line in lines if 'def rotation_q_wxyz(' not in line]; t='\n'.join(lines)+'\n'; t=t.replace('numpy.ndarray[tuple[typing.Literal[2], N], numpy.dtype[numpy.float64]]', 'numpy.ndarray[tuple[typing.Literal[2], typing.Any], numpy.dtype[numpy.float64]]'); p.write_text(t)"
+	python -c "from pathlib import Path; p=Path('python/rcs/_core/sim.pyi'); t=p.read_text(); t=t.replace('numpy.ndarray[tuple[typing.Literal[2], N], numpy.dtype[numpy.float64]]', 'numpy.ndarray[tuple[typing.Literal[2], typing.Any], numpy.dtype[numpy.float64]]'); t=t.replace(', max_buffer_frames: int = 100', ''); p.write_text(t)"
 	python scripts/generate_common_typing.py
 	ruff check --fix python/rcs/_core python/rcs/common_typing.py
 	isort python/rcs/_core python/rcs/common_typing.py

python/rcs/_core/common.pyi

@@ -197,7 +197,6 @@ class Pose:
     def pose_matrix(self) -> numpy.ndarray[tuple[typing.Literal[4], typing.Literal[4]], numpy.dtype[numpy.float64]]: ...
     def rotation_m(self) -> numpy.ndarray[tuple[typing.Literal[3], typing.Literal[3]], numpy.dtype[numpy.float64]]: ...
     def rotation_q(self) -> numpy.ndarray[tuple[typing.Literal[4]], numpy.dtype[numpy.float64]]: ...
-    def rotation_q_wxyz(self) -> numpy.ndarray[tuple[typing.Literal[4]], numpy.dtype[numpy.float64]]: ...
     def rotation_rpy(self) -> RPY: ...
     def rotvec(self) -> numpy.ndarray[tuple[typing.Literal[6]], numpy.dtype[numpy.float64]]: ...
     def total_angle(self) -> float: ...
@@ -242,7 +241,7 @@ class Robot:
 class RobotConfig:
     attachment_site: str
     dof: int
-    joint_limits: numpy.ndarray[tuple[typing.Literal[2], N], numpy.dtype[numpy.float64]]
+    joint_limits: numpy.ndarray[tuple[typing.Literal[2], typing.Any], numpy.dtype[numpy.float64]]
     kinematic_model_path: str
     q_home: numpy.ndarray | None
     robot_platform: RobotPlatform
@@ -252,7 +251,7 @@ class RobotConfig:
         self,
         robot_type: RobotType = ...,
         dof: int = 7,
-        joint_limits: numpy.ndarray[tuple[typing.Literal[2], N], numpy.dtype[numpy.float64]] = ...,
+        joint_limits: numpy.ndarray[tuple[typing.Literal[2], typing.Any], numpy.dtype[numpy.float64]] = ...,
         robot_platform: RobotPlatform = ...,
         tcp_offset: Pose = ...,
         attachment_site: str = "attachment_site",

python/rcs/_core/sim.pyi

@@ -109,9 +109,7 @@ class SimCameraConfig(rcs._core.common.BaseCameraConfig):
     ) -> None: ...
 
 class SimCameraSet:
-    def __init__(
-        self, sim: Sim, cameras: dict[str, SimCameraConfig], render_on_demand: bool = True, max_buffer_frames: int = 100
-    ) -> None: ...
+    def __init__(self, sim: Sim, cameras: dict[str, SimCameraConfig], render_on_demand: bool = True) -> None: ...
     def buffer_size(self) -> int: ...
     def clear_buffer(self) -> None: ...
     def get_latest_frameset(self) -> FrameSet | None: ...
@@ -201,7 +199,7 @@ class SimRobotConfig(rcs._core.common.RobotConfig):
     arm_collision_geoms: list[str]
     base: str
     dof: int
-    joint_limits: numpy.ndarray[tuple[typing.Literal[2], N], numpy.dtype[numpy.float64]]
+    joint_limits: numpy.ndarray[tuple[typing.Literal[2], typing.Any], numpy.dtype[numpy.float64]]
     joint_rotational_tolerance: float
     joints: list[str]
     seconds_between_callbacks: float
@@ -248,7 +246,7 @@ class SimRobotConfig(rcs._core.common.RobotConfig):
         ],
         base: str = "base",
         dof: int = 7,
-        joint_limits: numpy.ndarray[tuple[typing.Literal[2], N], numpy.dtype[numpy.float64]] = ...,
+        joint_limits: numpy.ndarray[tuple[typing.Literal[2], typing.Any], numpy.dtype[numpy.float64]] = ...,
     ) -> None: ...
     def add_prefix(self, id: str) -> None: ...
 

python/rcs/camera/sim.py

@@ -26,18 +26,12 @@ def __init__(
         cameras: dict[str, SimCameraConfig],
         physical_units: bool = False,
         render_on_demand: bool = True,
-        max_buffer_frames: int = 100,
     ):
         self._logger = logging.getLogger(__name__)
         self.cameras = cameras
         self.physical_units = physical_units
 
-        super().__init__(
-            simulation,
-            cameras,
-            render_on_demand=render_on_demand,
-            max_buffer_frames=max_buffer_frames,
-        )
+        super().__init__(simulation, cameras, render_on_demand=render_on_demand)
         self._sim: sim.Sim
 
     def get_latest_frames(self) -> FrameSet | None:

python/rcs/envs/tasks.py

@@ -130,8 +130,10 @@ def reset(
         pose_in_world_frame = self.center2world * pose_in_center_frame
 
         # qpos array format for a free joint: [x, y, z, qw, qx, qy, qz]
+        quat_xyzw = pose_in_world_frame.rotation_q()
+        quat_wxyz = np.array([quat_xyzw[3], quat_xyzw[0], quat_xyzw[1], quat_xyzw[2]])
         self.get_wrapper_attr("sim").data.joint(self.obj_joint_name).qpos = np.append(
-            pose_in_world_frame.translation(), pose_in_world_frame.rotation_q_wxyz()
+            pose_in_world_frame.translation(), quat_wxyz
         )
 
         # reset of remaining stack, must happen after our reset!

python/rcs/sim/composer.py

@@ -65,7 +65,8 @@ def _find_camera(self, name: str) -> Optional[mujoco._specs.MjsCamera]:
 
     def _apply_pose(self, body: mujoco._specs.MjsBody, pose: Pose):
         body.pos = list(pose.translation())
-        body.quat = list(pose.rotation_q_wxyz())
+        quat_xyzw = pose.rotation_q()
+        body.quat = [quat_xyzw[3], quat_xyzw[0], quat_xyzw[1], quat_xyzw[2]]
 
     def add_camera(
         self,