Merge pull request #17 from ami-iit/ramps_and_jumps

ami-iit · Jul 29, 2024 · f6ccb4f · f6ccb4f
2 parents 3cd1f2f + 22f07f2
commit f6ccb4f
Show file tree

Hide file tree

Showing 4 changed files with 1,276 additions and 0 deletions.
diff --git a/src/hippopt/robot_planning/utilities/terrain_descriptor.py b/src/hippopt/robot_planning/utilities/terrain_descriptor.py
@@ -9,6 +9,7 @@ class TerrainDescriptor(abc.ABC):
     _height_function: cs.Function | None = dataclasses.field(default=None)
     _normal_direction_function: cs.Function | None = dataclasses.field(default=None)
     _orientation_function: cs.Function | None = dataclasses.field(default=None)
+    _transform_function: cs.Function | None = dataclasses.field(default=None)
     _point_position_name: str = dataclasses.field(default="point_position")
     _name: str = dataclasses.field(default=None)
     _options: dict = dataclasses.field(default_factory=dict)
@@ -78,6 +79,25 @@ def create_orientation_function(self) -> cs.Function:
             self._options,
         )
 
+    def create_transform_function(self) -> cs.Function:
+        point_position = cs.MX.sym(self.get_point_position_name(), 3)
+        height = self.height_function()(point_position)
+        rotation = self.orientation_function()(point_position)
+
+        point_on_surface = cs.vertcat(point_position[:2], point_position[2] - height)
+        top_rows = cs.horzcat(rotation, point_on_surface)
+        fourth_row = cs.DM.eye(4)[3, :]
+        transform_matrix = cs.vertcat(top_rows, fourth_row)
+
+        return cs.Function(
+            "terrain_transform",
+            [point_position],
+            [transform_matrix],
+            [self.get_point_position_name()],
+            ["plane_transform"],
+            self._options,
+        )
+
     def height_function(self) -> cs.Function:
         if not isinstance(self._height_function, cs.Function):
             self._height_function = self.create_height_function()
@@ -96,6 +116,12 @@ def orientation_function(self) -> cs.Function:
 
         return self._orientation_function
 
+    def transform_function(self) -> cs.Function:
+        if not isinstance(self._transform_function, cs.Function):
+            self._transform_function = self.create_transform_function()
+
+        return self._transform_function
+
     def get_point_position_name(self) -> str:
         return self._point_position_name
 
@@ -109,3 +135,4 @@ def invalidate_functions(self) -> None:
         self._height_function = None
         self._normal_direction_function = None
         self._orientation_function = None
+        self._transform_function = None