Run pre-commit

airo-ugent · Sep 11, 2024 · 37c307e · 37c307e
1 parent 9182fdd
commit 37c307e
Show file tree

Hide file tree

Showing 6 changed files with 59 additions and 84 deletions.
diff --git a/airo_drake/building/manipulator.py b/airo_drake/building/manipulator.py
@@ -2,7 +2,7 @@
 import numpy as np
 from airo_typing import HomogeneousMatrixType
 from pydrake.math import RigidTransform, RollPitchYaw
-from pydrake.multibody.tree import ModelInstanceIndex, Frame
+from pydrake.multibody.tree import Frame, ModelInstanceIndex
 from pydrake.planning import RobotDiagramBuilder
 
 # X_URBASE_ROSBASE is the 180 rotation between ROS URDF base and the UR control box base
@@ -13,14 +13,14 @@
 
 
 def add_manipulator(
-        robot_diagram_builder: RobotDiagramBuilder,
-        name: str,
-        gripper_name: str,
-        arm_transform: HomogeneousMatrixType | None = None,
-        gripper_transform: HomogeneousMatrixType | None = None,
-        static_arm: bool = False,
-        static_gripper: bool = False,
-        parent_frame: Frame | None = None,
+    robot_diagram_builder: RobotDiagramBuilder,
+    name: str,
+    gripper_name: str,
+    arm_transform: HomogeneousMatrixType | None = None,
+    gripper_transform: HomogeneousMatrixType | None = None,
+    static_arm: bool = False,
+    static_gripper: bool = False,
+    parent_frame: Frame | None = None,
 ) -> tuple[ModelInstanceIndex, ModelInstanceIndex]:
     """Add a manipulator (a robot arm with a gripper) to the robot diagram builder.
     Looks up the URDF files for the robot and gripper and welds them together.
@@ -53,9 +53,7 @@ def add_manipulator(
     if static_arm:
         arm_urdf = airo_models.urdf.read_urdf(arm_urdf_path)
         airo_models.urdf.make_static(arm_urdf)
-        arm_urdf_path = airo_models.urdf.write_urdf_to_tempfile(
-            arm_urdf, arm_urdf_path, prefix=f"{name}_static_"
-        )
+        arm_urdf_path = airo_models.urdf.write_urdf_to_tempfile(arm_urdf, arm_urdf_path, prefix=f"{name}_static_")
 
     arm_index = parser.AddModels(arm_urdf_path)[0]
 

diff --git a/airo_drake/building/mobile_platform.py b/airo_drake/building/mobile_platform.py
@@ -1,4 +1,4 @@
-from typing import Tuple, List
+from typing import List, Tuple
 
 import airo_models
 import numpy as np
@@ -15,16 +15,20 @@
 
 
 def add_mobile_platform(
-        robot_diagram_builder: RobotDiagramBuilder,
-        drive_positions: Tuple[Vector2DType] = (np.array([1, -0.5]), np.array([1, 0.5]), np.array([-1, -0.5]),
-                                                np.array([-1, 0.5])),
-        battery_position: Vector2DType = np.array([0, 0.5]),
-        cpu_position: Vector2DType = np.array([0, -0.5]),
-        side_height: float = 0.43,
-        side_length: float = 0.69,
-        roof_width: float = 0.525,
-        roof_thickness: float = 0.03,
-        brick_size: float = 0.233
+    robot_diagram_builder: RobotDiagramBuilder,
+    drive_positions: Tuple[Vector2DType] = (
+        np.array([1, -0.5]),
+        np.array([1, 0.5]),
+        np.array([-1, -0.5]),
+        np.array([-1, 0.5]),
+    ),
+    battery_position: Vector2DType = np.array([0, 0.5]),
+    cpu_position: Vector2DType = np.array([0, -0.5]),
+    side_height: float = 0.43,
+    side_length: float = 0.69,
+    roof_width: float = 0.525,
+    roof_thickness: float = 0.03,
+    brick_size: float = 0.233,
 ) -> Tuple[ModelInstanceIndex, List[ModelInstanceIndex]]:
     """Add a mobile platform on wheels to the robot diagram builder.
 
@@ -58,8 +62,8 @@ def add_mobile_platform(
     mobi_urdf_path = airo_models.get_urdf_path("kelo_robile")
 
     drive_transforms = [
-        RigidTransform(p=[brick_size * p[0], brick_size * p[1], 0], rpy=RollPitchYaw([0, 0, 0])) for p in
-        drive_positions
+        RigidTransform(p=[brick_size * p[0], brick_size * p[1], 0], rpy=RollPitchYaw([0, 0, 0]))
+        for p in drive_positions
     ]
 
     mobi_model_index = parser.AddModels(mobi_urdf_path)[0]
@@ -75,15 +79,17 @@ def add_mobile_platform(
         plant.WeldFrames(robot_root_frame, brick_frame, drive_transform)
         mobi_part_indices.append(brick_index)
 
-    battery_transform = RigidTransform(p=[brick_size * battery_position[0], brick_size * battery_position[1], 0],
-                                       rpy=RollPitchYaw([0, 0, 0]))
+    battery_transform = RigidTransform(
+        p=[brick_size * battery_position[0], brick_size * battery_position[1], 0], rpy=RollPitchYaw([0, 0, 0])
+    )
     battery_index = parser.AddModels(airo_models.get_urdf_path("kelo_robile_battery"))[0]
     battery_frame = plant.GetFrameByName("base_link", battery_index)
     plant.WeldFrames(robot_root_frame, battery_frame, battery_transform)
     mobi_part_indices.append(battery_index)
 
-    cpu_transform = RigidTransform(p=[brick_size * cpu_position[0], brick_size * cpu_position[1], 0],
-                                   rpy=RollPitchYaw([0, 0, 0]))
+    cpu_transform = RigidTransform(
+        p=[brick_size * cpu_position[0], brick_size * cpu_position[1], 0], rpy=RollPitchYaw([0, 0, 0])
+    )
     cpu_index = parser.AddModels(airo_models.get_urdf_path("kelo_robile_cpu"))[0]
     cpu_frame = plant.GetFrameByName("base_link", cpu_index)
     plant.WeldFrames(robot_root_frame, cpu_frame, cpu_transform)
@@ -119,8 +125,9 @@ def add_mobile_platform(
     return mobi_model_index, mobi_part_indices
 
 
-def attach_mobile_platform_to_world(robot_diagram_builder: RobotDiagramBuilder, mobi_model_index: ModelInstanceIndex,
-                                    static_platform: bool):
+def attach_mobile_platform_to_world(
+    robot_diagram_builder: RobotDiagramBuilder, mobi_model_index: ModelInstanceIndex, static_platform: bool
+):
     """Attach the mobile platform to the world frame with a planar joint. This allows the mobile platform to
     move around by setting the (x, y, theta) values of the joint.
 

diff --git a/airo_drake/scene.py b/airo_drake/scene.py
@@ -1,11 +1,10 @@
 from dataclasses import dataclass
+from typing import List
 
 from pydrake.geometry import Meshcat
 from pydrake.multibody.tree import ModelInstanceIndex
 from pydrake.planning import RobotDiagram
 
-from typing import List
-
 
 @dataclass
 class MobilePlatformWithSingleArmScene:

diff --git a/airo_drake/time_parametrization/toppra.py b/airo_drake/time_parametrization/toppra.py
@@ -6,12 +6,12 @@
 
 
 def time_parametrize_toppra_mobile_platform(
-        plant: MultibodyPlant,
-        poses: JointPathType,  # TODO type
-        linear_velocity_limit: float = 1.0,
-        angular_velocity_limit: float = np.pi / 8,
-        linear_acceleration_limit: float = 0.25,
-        angular_acceleration_limit: float = 0.025
+    plant: MultibodyPlant,
+    poses: JointPathType,  # TODO type
+    linear_velocity_limit: float = 1.0,
+    angular_velocity_limit: float = np.pi / 8,
+    linear_acceleration_limit: float = 0.25,
+    angular_acceleration_limit: float = 0.025,
 ) -> PathParameterizedTrajectory:
     """Recalculate the timing of a path or trajectory to respect linear and angular velocity and acceleration limits
     using TOPP-RA.
@@ -36,7 +36,8 @@ def time_parametrize_toppra_mobile_platform(
     gridpoints = Toppra.CalcGridPoints(pose_trajectory, CalcGridPointsOptions())
 
     acceleration_limits_lower = np.array(
-        [-linear_acceleration_limit, -linear_acceleration_limit, -angular_acceleration_limit])
+        [-linear_acceleration_limit, -linear_acceleration_limit, -angular_acceleration_limit]
+    )
     acceleration_limits_upper = -acceleration_limits_lower
     velocity_limits_lower = np.array([-linear_velocity_limit, -linear_velocity_limit, -angular_velocity_limit])
     velocity_limits_upper = -velocity_limits_lower
@@ -55,10 +56,10 @@ def time_parametrize_toppra_mobile_platform(
 
 
 def time_parametrize_toppra(
-        plant: MultibodyPlant,
-        joints: JointPathType | Trajectory,
-        joint_speed_limit: float = 2.0,  # Max 180 degrees/s ~ 3.14 rad/s
-        joint_acceleration_limit: float = 4.0,  # UR recommends < 800 degrees/s^2 ~ 13.9 rad/s^2
+    plant: MultibodyPlant,
+    joints: JointPathType | Trajectory,
+    joint_speed_limit: float = 2.0,  # Max 180 degrees/s ~ 3.14 rad/s
+    joint_acceleration_limit: float = 4.0,  # UR recommends < 800 degrees/s^2 ~ 13.9 rad/s^2
 ) -> PathParameterizedTrajectory:
     """Recalculate the timing of a path or trajectory to respect joint speed and acceleration limits using TOPP-RA.
 

diff --git a/airo_drake/visualization/mobile_platform.py b/airo_drake/visualization/mobile_platform.py
@@ -53,4 +53,4 @@ def animate_mobile_platform_trajectory(
         robot_diagram.ForcedPublish(context)
 
     meshcat.StopRecording()
-    meshcat.PublishRecording()
+    meshcat.PublishRecording()
diff --git a/notebooks/04_mobile_platform.ipynb b/notebooks/04_mobile_platform.ipynb
@@ -27,17 +27,9 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "INFO:drake:Meshcat listening for connections at http://localhost:7000\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "import numpy as np\n",
     "from pydrake.planning import RobotDiagramBuilder\n",
@@ -66,7 +58,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -87,7 +79,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -96,20 +88,9 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "MobilePlatformWithSingleArmScene(robot_diagram=<pydrake.planning.RobotDiagram object at 0x7fa3c0924370>, mobile_platform_index=ModelInstanceIndex(2), arm_index=ModelInstanceIndex(12), gripper_index=ModelInstanceIndex(13), meshcat=<pydrake.geometry.Meshcat object at 0x7fa3bbb1cdb0>)"
-      ]
-     },
-     "execution_count": 4,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
+   "outputs": [],
    "source": [
     "robot_diagram, context = finish_build(robot_diagram_builder)\n",
     "\n",
@@ -136,20 +117,9 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "9"
-      ]
-     },
-     "execution_count": 5,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
+   "outputs": [],
    "source": [
     "diagram = scene.robot_diagram\n",
     "plant = diagram.plant()\n",
@@ -161,7 +131,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [