Refactored MotionPlanner so that objective can be passed separately. (#…

…277)

* Refactored MotionPlanner so that objective can be passed separately.

* Made MotionPlannerObjective accessible from theseus.utils.examples.

* Removed irrelevant text from notebooks.

examples/motion_planning_2d.py

@@ -69,15 +69,15 @@ def run_learning_loop(cfg):
     )
 
     motion_planner = theg.MotionPlanner(
-        cfg.img_size,
-        cfg.obs_params.safety_dist + cfg.robot_radius,
-        cfg.total_time,
-        cfg.obs_params.weight,
-        cfg.gp_params.Qc_inv,
-        cfg.num_time_steps,
         cfg.optim_params.method,
         cfg.optim_params.max_iters,
-        cfg.optim_params.step_size,
+        step_size=cfg.optim_params.step_size,
+        map_size=cfg.img_size,
+        epsilon_dist=cfg.obs_params.safety_dist + cfg.robot_radius,
+        total_time=cfg.total_time,
+        collision_weight=cfg.obs_params.weight,
+        Qc_inv=cfg.gp_params.Qc_inv,
+        num_time_steps=cfg.num_time_steps,
         use_single_collision_weight=True,
         device=cfg.device,
     )

theseus/utils/examples/__init__.py

@@ -10,6 +10,7 @@
     from .motion_planning import (
         InitialTrajectoryModel,
         MotionPlanner,
+        MotionPlannerObjective,
         ScalarCollisionWeightAndCostEpstModel,
         ScalarCollisionWeightModel,
         TrajectoryDataset,

theseus/utils/examples/motion_planning/__init__.py

@@ -9,4 +9,4 @@
     ScalarCollisionWeightAndCostEpstModel,
     ScalarCollisionWeightModel,
 )
-from .motion_planner import MotionPlanner
+from .motion_planner import MotionPlanner, MotionPlannerObjective

theseus/utils/examples/motion_planning/models.py

@@ -121,7 +121,7 @@ def __init__(
             nn.ReLU(),
             nn.Linear(hid_size, hid_size),
             nn.ReLU(),
-            nn.Linear(hid_size, 4 * (planner.num_time_steps + 1)),
+            nn.Linear(hid_size, 4 * (planner.objective.num_time_steps + 1)),
         )
 
         # Learns a quadratic offset in normal direction to bend the mean trajectory.
@@ -141,7 +141,7 @@ def init_weights(m_):
 
         self.bend_factor.apply(init_weights)
 
-        self.dt = planner.total_time / planner.num_time_steps
+        self.dt = planner.objective.total_time / planner.objective.num_time_steps
 
         self.num_images = max_num_images
 
@@ -157,7 +157,7 @@ def forward(self, batch: Dict[str, Any]):
             one_hot_dummy[batch_idx, idx] = 1
 
         # Compute straight line positions to use as mean of initial trajectory
-        trajectory_len = self.aux_motion_planner.trajectory_len
+        trajectory_len = self.aux_motion_planner.objective.trajectory_len
         dist_vec = goal - start
         pos_incr_per_step = dist_vec / (trajectory_len - 1)
         trajectory = torch.zeros(start.shape[0], 4 * trajectory_len).to(

theseus/utils/examples/motion_planning/motion_planner.py

@@ -11,7 +11,7 @@
 import theseus as th
 
 
-class MotionPlanner:
+class MotionPlannerObjective(th.Objective):
     def __init__(
         self,
         map_size: int,
@@ -20,25 +20,27 @@ def __init__(
         collision_weight: float,
         Qc_inv: List[List[int]],
         num_time_steps: int,
-        optim_method: str,
-        max_optim_iters: int,
-        step_size: float = 1.0,
         use_single_collision_weight: bool = True,
-        device: str = "cpu",
         dtype: torch.dtype = torch.double,
     ):
+        for v in [
+            map_size,
+            epsilon_dist,
+            total_time,
+            collision_weight,
+            Qc_inv,
+            num_time_steps,
+        ]:
+            assert v is not None
+
+        super().__init__(dtype=dtype)
         self.map_size = map_size
         self.epsilon_dist = epsilon_dist
         self.total_time = total_time
         self.collision_weight = collision_weight
         self.Qc_inv = copy.deepcopy(Qc_inv)
         self.num_time_steps = num_time_steps
-        self.optim_method = optim_method
-        self.max_optim_iters = max_optim_iters
-        self.step_size = step_size
         self.use_single_collision_weight = use_single_collision_weight
-        self.device = device
-        self.dtype = dtype
 
         self.trajectory_len = num_time_steps + 1
 
@@ -112,27 +114,24 @@ def __init__(
         # --------------------------------------------------------------------------- #
         # ------------------------------ Cost functions ----------------------------- #
         # --------------------------------------------------------------------------- #
-        # Create a Theseus objective for adding the cost functions
-        objective = th.Objective(dtype=self.dtype)
-
         # First create the cost functions for the end point positions and velocities
         # which are hard constraints, and can be implemented via Difference cost
         # functions.
-        objective.add(
+        self.add(
             th.Difference(poses[0], start_point, boundary_cost_weight, name="pose_0")
         )
-        objective.add(
+        self.add(
             th.Difference(
                 velocities[0],
                 th.Point2(tensor=torch.zeros(1, 2, dtype=dtype)),
                 boundary_cost_weight,
                 name="vel_0",
             )
         )
-        objective.add(
+        self.add(
             th.Difference(poses[-1], goal_point, boundary_cost_weight, name="pose_N")
         )
-        objective.add(
+        self.add(
             th.Difference(
                 velocities[-1],
                 th.Point2(tensor=torch.zeros(1, 2, dtype=dtype)),
@@ -145,7 +144,7 @@ def __init__(
         # cost weights created above. We need a separate cost function for each time
         # step
         for i in range(1, self.trajectory_len):
-            objective.add(
+            self.add(
                 th.eb.Collision2D(
                     poses[i],
                     sdf_origin,
@@ -158,7 +157,7 @@ def __init__(
                     name=f"collision_{i}",
                 )
             )
-            objective.add(
+            self.add(
                 (
                     th.eb.GPMotionModel(
                         poses[i - 1],
@@ -172,25 +171,64 @@ def __init__(
                 )
             )
 
+
+class MotionPlanner:
+    # If objective is given, this overrides problem arguments
+    def __init__(
+        self,
+        optim_method: str,
+        max_optim_iters: int,
+        step_size: float = 1.0,
+        objective: Optional[MotionPlannerObjective] = None,
+        device: str = "cpu",
+        dtype: torch.dtype = torch.double,
+        # The following are only used if objective is None
+        map_size: Optional[int] = None,
+        epsilon_dist: Optional[float] = None,
+        total_time: Optional[float] = None,
+        collision_weight: Optional[float] = None,
+        Qc_inv: Optional[List[List[int]]] = None,
+        num_time_steps: Optional[int] = None,
+        use_single_collision_weight: bool = True,
+    ):
+        if objective is None:
+            self.objective = MotionPlannerObjective(
+                map_size,
+                epsilon_dist,
+                total_time,
+                collision_weight,
+                Qc_inv,
+                num_time_steps,
+                use_single_collision_weight=use_single_collision_weight,
+                dtype=dtype,
+            )
+        else:
+            self.objective = objective
+
+        self.optim_method = optim_method
+        self.max_optim_iters = max_optim_iters
+        self.step_size = step_size
+        self.device = device
+        self.dtype = dtype
+
         # Finally, create the Nonlinear Least Squares optimizer for this objective
         # and wrap both into a TheseusLayer
         optimizer: th.NonlinearLeastSquares
         if optim_method == "gauss_newton":
             optimizer = th.GaussNewton(
-                objective,
+                self.objective,
                 th.CholeskyDenseSolver,
                 max_iterations=max_optim_iters,
                 step_size=step_size,
             )
         elif optim_method == "levenberg_marquardt":
             optimizer = th.LevenbergMarquardt(
-                objective,
+                self.objective,
                 th.CholeskyDenseSolver,
                 max_iterations=max_optim_iters,
                 step_size=step_size,
             )
 
-        self.objective = objective
         self.layer = th.TheseusLayer(optimizer)
         self.layer.to(device=device, dtype=dtype)
 
@@ -249,10 +287,10 @@ def get_variable_values_from_straight_line(
         # Returns a dictionary of variable names to values that represent a straight
         # line trajectory from start to goal.
         start_goal_dist = goal - start
-        avg_vel = start_goal_dist / self.total_time
-        unit_trajectory_len = start_goal_dist / (self.trajectory_len - 1)
+        avg_vel = start_goal_dist / self.objective.total_time
+        unit_trajectory_len = start_goal_dist / (self.objective.trajectory_len - 1)
         input_dict: Dict[str, torch.Tensor] = {}
-        for i in range(self.trajectory_len):
+        for i in range(self.objective.trajectory_len):
             input_dict[f"pose_{i}"] = start + unit_trajectory_len * i
             input_dict[f"vel_{i}"] = avg_vel
         return input_dict
@@ -262,7 +300,7 @@ def get_random_variable_values(
     ) -> Dict[str, torch.Tensor]:
         # Returns a dictionary of variable names with random initial poses.
         input_dict: Dict[str, torch.Tensor] = {}
-        for i in range(self.trajectory_len):
+        for i in range(self.objective.trajectory_len):
             input_dict[f"pose_{i}"] = torch.randn_like(start)
             input_dict[f"vel_{i}"] = torch.randn_like(start)
         return input_dict
@@ -273,9 +311,9 @@ def get_variable_values_from_trajectory(
         # Returns a dictionary of variable names to values, so that values
         # are assigned with the data from the given trajectory. Trajectory should be a
         # tensor of shape (batch_size, 4, planner.trajectory_len).
-        assert trajectory.shape[1:] == (4, self.trajectory_len)
+        assert trajectory.shape[1:] == (4, self.objective.trajectory_len)
         input_dict: Dict[str, torch.Tensor] = {}
-        for i in range(self.trajectory_len):
+        for i in range(self.objective.trajectory_len):
             input_dict[f"pose_{i}"] = trajectory[:, :2, i]
             input_dict[f"vel_{i}"] = trajectory[:, :2, i]
         return input_dict
@@ -296,11 +334,11 @@ def get_trajectory(
         trajectory = torch.empty(
             self.objective.batch_size,
             4,
-            self.trajectory_len,
+            self.objective.trajectory_len,
             device=self.objective.device,
         )
         variables = self.objective.optim_vars
-        for i in range(self.trajectory_len):
+        for i in range(self.objective.trajectory_len):
             if values_dict is None:
                 trajectory[:, :2, i] = variables[f"pose_{i}"].tensor.clone()
                 trajectory[:, 2:, i] = variables[f"vel_{i}"].tensor.clone()
@@ -321,16 +359,16 @@ def get_total_squared_errors(self) -> Tuple[torch.Tensor, torch.Tensor]:
 
     def copy(self, collision_weight: Optional[float] = None) -> "MotionPlanner":
         return MotionPlanner(
-            self.map_size,
-            self.epsilon_dist,
-            self.total_time,
-            collision_weight or self.collision_weight,
-            self.Qc_inv,
-            self.num_time_steps,
             self.optim_method,
             self.max_optim_iters,
             step_size=self.step_size,
-            use_single_collision_weight=self.use_single_collision_weight,
+            map_size=self.objective.map_size,
+            epsilon_dist=self.objective.epsilon_dist,
+            total_time=self.objective.total_time,
+            collision_weight=collision_weight or self.objective.collision_weight,
+            Qc_inv=self.objective.Qc_inv,
+            num_time_steps=self.objective.num_time_steps,
+            use_single_collision_weight=self.objective.use_single_collision_weight,
             device=self.device,
             dtype=self.dtype,
         )