From a45dc996f6218772ff0049ae85245b8c78863fc8 Mon Sep 17 00:00:00 2001
From: victorlouisdg <victorlouisdg@gmail.com>
Date: Thu, 21 Mar 2024 12:44:25 +0100
Subject: [PATCH] add_wall() scene building utility

---
 airo_drake/__init__.py      |  2 ++
 airo_drake/building/wall.py | 60 +++++++++++++++++++++++++++++++++++++
 airo_drake/scene.py         |  8 ++---
 3 files changed, 66 insertions(+), 4 deletions(-)
 create mode 100644 airo_drake/building/wall.py

diff --git a/airo_drake/__init__.py b/airo_drake/__init__.py
index c09f7f2..a64e0d1 100644
--- a/airo_drake/__init__.py
+++ b/airo_drake/__init__.py
@@ -7,6 +7,7 @@
 from airo_drake.building.floor import add_floor
 from airo_drake.building.manipulator import X_URBASE_ROSBASE, X_URTOOL0_ROBOTIQ, add_manipulator
 from airo_drake.building.meshcat import add_meshcat
+from airo_drake.building.wall import add_wall
 from airo_drake.path.analysis import (
     calculate_joint_path_distances,
     calculate_joint_path_length,
@@ -28,6 +29,7 @@
 
 __all__ = [
     "add_floor",
+    "add_wall",
     "add_manipulator",
     "X_URBASE_ROSBASE",
     "X_URTOOL0_ROBOTIQ",
diff --git a/airo_drake/building/wall.py b/airo_drake/building/wall.py
new file mode 100644
index 0000000..6cbd2f0
--- /dev/null
+++ b/airo_drake/building/wall.py
@@ -0,0 +1,60 @@
+import airo_models
+import numpy as np
+from airo_typing import Vector3DType
+from pydrake.math import RigidTransform
+from pydrake.multibody.tree import ModelInstanceIndex
+from pydrake.planning import RobotDiagramBuilder
+
+
+def add_wall(
+    robot_diagram_builder: RobotDiagramBuilder,
+    type: str = "XZ",
+    height: float = 2.0,
+    width: float = 2.5,
+    thickness: float = 0.2,
+    position: Vector3DType = np.zeros(3),
+    name: str = "wall",
+) -> ModelInstanceIndex:
+    """Adds a wall to the robot diagram builder.
+
+    The wall is represented as a box with the specified dimensions, type, and position.
+    The wall is automatically welded to the world frame for a static setup.
+
+    Args:
+        robot_diagram_builder: The robot diagram builder to which the wall will be added.
+        type: The orientation of the wall.  Must be 'XZ' (parallel to the X-Z plane)
+              or 'YZ' (parallel to the Y-Z plane).
+        height: The height of the wall.
+        width: The width of the wall (relevant for its orientation).
+        thickness: The thickness of the wall (could be relevant for collision checking).
+        position: The 3D position (x, y, z) of the center of the wall.
+        name: The name of the wall (for easy reference).
+
+    Returns:
+        The ModelInstanceIndex of the added wall.
+    """
+    plant = robot_diagram_builder.plant()
+    parser = robot_diagram_builder.parser()
+    parser.SetAutoRenaming(True)
+
+    z_size = height
+
+    if type == "XZ":
+        x_size = width
+        y_size = thickness
+    elif type == "YZ":
+        x_size = thickness
+        y_size = width
+    else:
+        raise ValueError("Invalid wall type, must be 'XZ' or 'YZ'")
+
+    wall_urdf_path = airo_models.box_urdf_path((x_size, y_size, z_size), name)
+    wall_index = parser.AddModels(wall_urdf_path)[0]
+
+    wall_transform = RigidTransform(p=position)
+    world_frame = plant.world_frame()
+    wall_frame = plant.GetFrameByName("base_link", wall_index)
+
+    plant.WeldFrames(world_frame, wall_frame, wall_transform)
+
+    return wall_index
diff --git a/airo_drake/scene.py b/airo_drake/scene.py
index dd6a903..f117bb8 100644
--- a/airo_drake/scene.py
+++ b/airo_drake/scene.py
@@ -20,8 +20,8 @@ class DualArmScene:
     """The most important objects when using Drake with two robot arms."""
 
     robot_diagram: RobotDiagram
-    arm_index_left: ModelInstanceIndex
-    arm_index_right: ModelInstanceIndex
-    gripper_index_left: ModelInstanceIndex | None = None
-    gripper_index_right: ModelInstanceIndex | None = None
+    arm_left_index: ModelInstanceIndex
+    arm_right_index: ModelInstanceIndex
+    gripper_left_index: ModelInstanceIndex | None = None
+    gripper_right_index: ModelInstanceIndex | None = None
     meshcat: Meshcat | None = None