From 729a52c4b7166d75ddb30f40c78d87db48c5863d Mon Sep 17 00:00:00 2001
From: NishanthJKumar <njk@mit.edu>
Date: Fri, 28 Jul 2023 08:45:55 -0400
Subject: [PATCH] changes for hacking; need to drill hole and modify platform

---
 predicators/ground_truth_models/spot/nsrts.py |  2 +-
 predicators/spot_utils/spot_utils.py          | 26 ++++++++++++++++---
 2 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/predicators/ground_truth_models/spot/nsrts.py b/predicators/ground_truth_models/spot/nsrts.py
index d3289c2460..30622d87bb 100644
--- a/predicators/ground_truth_models/spot/nsrts.py
+++ b/predicators/ground_truth_models/spot/nsrts.py
@@ -75,7 +75,7 @@ def _grasp_sampler(spot_interface: _SpotInterface, state: State,
     if objs[1].type.name == "bag":  # pragma: no cover
         return np.array([0.0, 0.0, 0.0, -1.0])
     if objs[1].type.name == "platform":  # pragma: no cover
-        return np.array([0.0, 0.0, 0.0, -1.0])
+        return np.array([0.0, 0.0, 0.0, 2.0])
     if objs[2].name == "low_wall_rack":  # pragma: no cover
         return np.array([0.0, 0.0, 0.1, 0.0])
     return np.array([0.0, 0.0, 0.0, 0.0])
diff --git a/predicators/spot_utils/spot_utils.py b/predicators/spot_utils/spot_utils.py
index 56360f0db0..7f49a22944 100644
--- a/predicators/spot_utils/spot_utils.py
+++ b/predicators/spot_utils/spot_utils.py
@@ -209,6 +209,7 @@ def __init__(self) -> None:
         self._force_horizontal_grasp = False
         self._force_squeeze_grasp = False
         self._force_top_down_grasp = False
+        self._force_forward_grasp = False
         self._image_source = "hand_color_image"
 
         self.hand_x, self.hand_y, self.hand_z = (0.80, 0, 0.45)
@@ -567,7 +568,7 @@ def params_spaces(self) -> Dict[str, Box]:
         """The parameter spaces for each of the controllers."""
         return {
             "navigate": Box(-5.0, 5.0, (3, )),
-            "grasp": Box(-1.0, 1.0, (4, )),
+            "grasp": Box(-1.0, 2.0, (4, )),
             "placeOnTop": Box(-5.0, 5.0, (3, )),
             # TODO: make this reasonable later, and remember for now these are absolute coordinates!
             "drag": Box(-45.0, 45.0, (2, )),
@@ -708,13 +709,20 @@ def graspController(self, objs: Sequence[Object], params: Array) -> None:
         """
         print("Grasp", objs)
         assert len(params) == 4
-        assert params[3] in [0, 1, -1]
+        assert params[3] in [0, 1, -1, 2]
         if params[3] == 1:
             self._force_horizontal_grasp = False
             self._force_top_down_grasp = True
+            self._force_forward_grasp = False
         elif params[3] == -1:
             self._force_horizontal_grasp = True
             self._force_top_down_grasp = False
+            self._force_forward_grasp = False
+        elif params[3] == 2:
+            self._force_horizontal_grasp = False
+            self._force_top_down_grasp = False
+            self._force_forward_grasp = True
+
         self.arm_object_grasp(objs[1])
         if not np.allclose(params[:3], [0.0, 0.0, 0.0]):
             self.hand_movement(params[:3], open_gripper=False)
@@ -850,7 +858,7 @@ def add_grasp_constraint(
 
         # For these options, we'll use a vector alignment constraint.
         use_vector_constraint = self._force_top_down_grasp or \
-            self._force_horizontal_grasp
+            self._force_horizontal_grasp or self._force_forward_grasp
 
         # Specify the frame we're using.
         grasp.grasp_params.grasp_params_frame_name = VISION_FRAME_NAME
@@ -880,6 +888,18 @@ def add_grasp_constraint(
                 # The axis in the vision frame is the positive z-axis
                 axis_to_align_with_ewrt_vo = geometry_pb2.Vec3(x=0, y=0, z=1)
 
+            if self._force_forward_grasp:
+                # Add a constraint that requests that the z-axis of the gripper
+                # is pointing in the positive z direction in the vision frame.
+                # This means that the gripper is constrained to be flat, as
+                # it usually is in the standard sow position.
+
+                # The axis on the gripper is the z-axis.
+                axis_on_gripper_ewrt_gripper = geometry_pb2.Vec3(x=0, y=0, z=1)
+
+                # The axis in the vision frame is the positive z-axis
+                axis_to_align_with_ewrt_vo = geometry_pb2.Vec3(x=0, y=0, z=1)
+
             # Add the vector constraint to our proto.
             constraint = grasp.grasp_params.allowable_orientation.add()
             constraint.vector_alignment_with_tolerance.\