Coverage for src/pyroboplan/trajectory/trajectory_optimization.py: 99%

190 statements  

« prev     ^ index     » next       coverage.py v7.6.12, created at 2025-03-15 22:28 -0400

1"""Utilities for trajectory optimization based planning.""" 

2 

3from functools import partial 

4import numpy as np 

5import warnings 

6 

7import pinocchio 

8from pydrake.autodiffutils import InitializeAutoDiff, ExtractGradient, ExtractValue 

9from pydrake.solvers import ( 

10 MathematicalProgram, 

11 QuadraticConstraint, 

12 Solve, 

13 SolverOptions, 

14 SnoptSolver, 

15) 

16 

17from pyroboplan.core.utils import ( 

18 calculate_collision_vector_and_jacobians, 

19 get_collision_pair_indices_from_bodies, 

20) 

21from pyroboplan.trajectory.polynomial import CubicPolynomialTrajectory 

22 

23 

24class CubicTrajectoryOptimizationOptions: 

25 """Options for cubic polynomial trajectory optimization.""" 

26 

27 def __init__( 

28 self, 

29 num_waypoints=3, 

30 samples_per_segment=11, 

31 min_segment_time=0.01, 

32 max_segment_time=10.0, 

33 min_vel=-np.inf, 

34 max_vel=np.inf, 

35 min_accel=-np.inf, 

36 max_accel=np.inf, 

37 min_jerk=-np.inf, 

38 max_jerk=np.inf, 

39 max_planning_time=30.0, 

40 check_collisions=False, 

41 collision_link_list=[], 

42 min_collision_dist=0.0, 

43 collision_influence_dist=0.1, 

44 collision_avoidance_cost_weight=0.0, 

45 ): 

46 """ 

47 Initializes a set of options for cubic polynomial trajectory optimization. 

48 

49 Parameters 

50 ---------- 

51 num_waypoints : int 

52 The number of waypoints in the trajectory. Must be greater than or equal to 2. 

53 samples_per_segment : int 

54 The number of samples to take along each trajectory segment for setting kinematic constraints. 

55 min_segment_time : float 

56 The minimum duration of a trajectory segment, in seconds. 

57 max_segment_time : float 

58 The maximum duration of a trajectory segment, in seconds. 

59 min_vel : float, or array-like 

60 The minimum velocity along the trajectory. 

61 If scalar, applies to all degrees of freedom; otherwise allows for different limits per degree of freedom. 

62 max_vel : float or array-like 

63 The maximum velocity along the trajectory. 

64 If scalar, applies to all degrees of freedom; otherwise allows for different limits per degree of freedom. 

65 min_accel : float, or array-like 

66 The minimum acceleration along the trajectory. 

67 If scalar, applies to all degrees of freedom; otherwise allows for different limits per degree of freedom. 

68 max_accel : float or array-like 

69 The maximum acceleration along the trajectory. 

70 If scalar, applies to all degrees of freedom; otherwise allows for different limits per degree of freedom. 

71 min_jerk : float, or array-like 

72 The minimum jerk along the trajectory. 

73 If scalar, applies to all degrees of freedom; otherwise allows for different limits per degree of freedom. 

74 max_jerk : float or array-like 

75 The maximum jerk along the trajectory. 

76 If scalar, applies to all degrees of freedom; otherwise allows for different limits per degree of freedom. 

77 max_planning_time : float 

78 The maximum planning time, in seconds. 

79 check_collisions : bool 

80 If true, adds collision constraints to trajectory optimization. 

81 collision_link_list : list[str] 

82 A list of collision objects over which to check minimum distance constraints. 

83 min_collision_dist : float 

84 The minimum allowable collision distance, in meters. 

85 collision_influence_dist : float 

86 The distance for collision/distance checks, in meters, above which to ignore results. 

87 collision_avoidance_cost_weight : float 

88 Weight for a cost function that penalizes the minimum distance from collision. 

89 If this is set to zero, the cost function is not added to the mathematical program. 

90 """ 

91 if num_waypoints < 2: 

92 raise ValueError( 

93 "The number of waypoints must be greater than or equal to 2." 

94 ) 

95 if min_segment_time <= 0: 

96 raise ValueError("The minimum segment time must be positive.") 

97 

98 self.num_waypoints = num_waypoints 

99 self.samples_per_segment = samples_per_segment 

100 self.min_segment_time = min_segment_time 

101 self.max_segment_time = max_segment_time 

102 self.min_vel = min_vel 

103 self.max_vel = max_vel 

104 self.min_accel = min_accel 

105 self.max_accel = max_accel 

106 self.min_jerk = min_jerk 

107 self.max_jerk = max_jerk 

108 

109 self.max_planning_time = max_planning_time 

110 

111 self.check_collisions = check_collisions 

112 self.collision_link_list = collision_link_list 

113 self.min_collision_dist = min_collision_dist 

114 self.collision_influence_dist = collision_influence_dist 

115 self.collision_avoidance_cost_weight = collision_avoidance_cost_weight 

116 

117 

118class CubicTrajectoryOptimization: 

119 """ 

120 Trajectory optimization based planner. 

121 

122 This uses the direct collocation approach to optimize over waypoint and collocation point placements that describe a multi-segment cubic polynomial trajectory. 

123 

124 Some good resources include: 

125 * Matthew Kelly's tutorials: https://www.matthewpeterkelly.com/tutorials/trajectoryOptimization/index.html 

126 * Russ Tedrake's manipulation course book: https://underactuated.mit.edu/trajopt.html 

127 """ 

128 

129 def __init__( 

130 self, model, collision_model, options=CubicTrajectoryOptimizationOptions() 

131 ): 

132 """ 

133 Creates an instance of a cubic trajectory optimization planner. 

134 

135 Parameters 

136 ---------- 

137 model : `pinocchio.Model` 

138 The model to use for this solver. 

139 collision_model : `pinocchio.GeometryModel` 

140 The model to use for collision checking. 

141 options : `CubicTrajectoryOptimizationOptions`, optional 

142 The options to use for planning. If not specified, default options are used. 

143 """ 

144 self.model = model 

145 self.collision_model = collision_model 

146 self.data = self.model.createData() 

147 self.collision_data = self.collision_model.createData() 

148 self.options = options 

149 

150 def _process_limits(self, limits, num_dofs, name): 

151 """ 

152 Helper function to process kinematics limits options. 

153 

154 * If the input limits are scalar, reshape them to the number of degrees of freedom. 

155 * Whether the input limits are scalar, a list, or a numpy array, always return a numpy array. 

156 * If the input limits are not scalar, but the wrong size, this will raise an Exception. 

157 

158 Parameters 

159 ---------- 

160 limits : None, float, or array-like 

161 The input limits in various formats. 

162 num_dofs : int 

163 The number of degrees of freedom. 

164 name : str 

165 The name of the input limits, to generate a descriptive error. 

166 

167 Return 

168 ------ 

169 numpy.ndarray 

170 The processed limits, for compatibility with trajectory optimization. 

171 """ 

172 limits = np.array(limits) 

173 if len(limits.shape) == 0 or limits.shape[0] == 1: 

174 limits = limits * np.ones(num_dofs) 

175 elif limits.shape != (num_dofs,): 175 ↛ 177line 175 didn't jump to line 177 because the condition on line 175 was always true

176 raise ValueError(f"{name} vector must have shape ({num_dofs},)") 

177 return limits 

178 

179 def _eval_position(self, x, x_d, xc_d, h, k, n, step): 

180 """ 

181 Helper function to symbolically evaluate a trajectory position. 

182 

183 This directly relates to equation 4.13 in https://epubs.siam.org/doi/10.1137/16M1062569. 

184 

185 Parameters 

186 ---------- 

187 x : pydrake.autodiffutils.AutoDiffXd 

188 The waypoint position points. 

189 x_d : pydrake.autodiffutils.AutoDiffXd 

190 The waypoint velocity points. 

191 xc_d : pydrake.autodiffutils.AutoDiffXd 

192 The collocation velocity points. 

193 h : pydrake.autodiffutils.AutoDiffXd 

194 The time segment durations. 

195 k : int 

196 The trajectory segment index. 

197 n : int 

198 The degree of freedom index. 

199 step : float 

200 The normalized distance, between 0 and 1, along that segment. 

201 

202 Return 

203 ------ 

204 pydrake.autodiffutils.AutodiffXd 

205 The position along the specified segment evaluated at the specified step. 

206 """ 

207 return ( 

208 x[k, n] 

209 + x_d[k, n] * (step * h[k]) 

210 + 0.5 

211 * (-3.0 * x_d[k, n] + 4.0 * xc_d[k, n] - x_d[k + 1, n]) 

212 * (step * h[k]) ** 2 

213 / h[k] 

214 + (2.0 / 3.0) 

215 * (x_d[k, n] - 2.0 * xc_d[k, n] + x_d[k + 1, n]) 

216 * (step * h[k]) ** 3 

217 / h[k] ** 2 

218 ) 

219 

220 def _eval_velocity(self, x_d, xc_d, h, k, n, step): 

221 """ 

222 Helper function to symbolically evaluate a trajectory velocity. 

223 

224 This is the first derivative of equation 4.13 in https://epubs.siam.org/doi/10.1137/16M1062569. 

225 

226 Parameters 

227 ---------- 

228 x_d : pydrake.autodiffutils.AutoDiffXd 

229 The waypoint velocity points. 

230 xc_d : pydrake.autodiffutils.AutoDiffXd 

231 The collocation velocity points. 

232 h : pydrake.autodiffutils.AutoDiffXd 

233 The time segment durations. 

234 k : int 

235 The trajectory segment index. 

236 n : int 

237 The degree of freedom index. 

238 step : float 

239 The normalized distance, between 0 and 1, along that segment. 

240 

241 Return 

242 ------ 

243 pydrake.autodiffutils.AutoDiffXd 

244 The velocity along the specified segment evaluated at the specified step. 

245 """ 

246 return ( 

247 x_d[k, n] 

248 + (-3.0 * x_d[k, n] + 4.0 * xc_d[k, n] - x_d[k + 1, n]) 

249 * (step * h[k]) 

250 / h[k] 

251 + 2.0 

252 * (x_d[k, n] - 2.0 * xc_d[k, n] + x_d[k + 1, n]) 

253 * (step * h[k]) ** 2 

254 / h[k] ** 2 

255 ) 

256 

257 def _eval_acceleration(self, x_d, xc_d, h, k, n, step): 

258 """ 

259 Helper function to symbolically evaluate a trajectory acceleration. 

260 

261 This is the second derivative of equation 4.13 in https://epubs.siam.org/doi/10.1137/16M1062569. 

262 

263 Parameters 

264 ---------- 

265 x_d : pydrake.autodiffutils.AutoDiffXd 

266 The waypoint velocity points. 

267 xc_d : pydrake.autodiffutils.AutoDiffXd 

268 The collocation velocity points. 

269 h : pydrake.autodiffutils.AutoDiffXd 

270 The time segment durations. 

271 k : int 

272 The trajectory segment index. 

273 n : int 

274 The degree of freedom index. 

275 step : float 

276 The normalized distance, between 0 and 1, along that segment. 

277 

278 Return 

279 ------ 

280 pydrake.autodiffutils.AutoDiffXd 

281 The acceleration along the specified segment evaluated at the specified step. 

282 """ 

283 return (-3.0 * x_d[k, n] + 4.0 * xc_d[k, n] - x_d[k + 1, n]) / h[k] + 4.0 * ( 

284 x_d[k, n] - 2.0 * xc_d[k, n] + x_d[k + 1, n] 

285 ) * (step * h[k]) / h[k] ** 2 

286 

287 def _eval_jerk(self, x_d, xc_d, h, k, n, step): 

288 """ 

289 Helper function to symbolically evaluate a trajectory jerk. 

290 

291 This is the third derivative of equation 4.13 in https://epubs.siam.org/doi/10.1137/16M1062569. 

292 

293 Parameters 

294 ---------- 

295 x_d : pydrake.autodiffutils.AutoDiffXd 

296 The waypoint velocity points. 

297 xc_d : pydrake.autodiffutils.AutoDiffXd 

298 The collocation velocity points. 

299 h : pydrake.autodiffutils.AutoDiffXd 

300 The time segment durations. 

301 k : int 

302 The trajectory segment index. 

303 n : int 

304 The degree of freedom index. 

305 step : float 

306 The normalized distance, between 0 and 1, along that segment. 

307 

308 Return 

309 ------ 

310 pydrake.autodiffutils.AutoDiffXd 

311 The jerk along the specified segment evaluated at the specified step. 

312 """ 

313 return 8.0 * (x_d[k, n] - 2.0 * xc_d[k, n] + x_d[k + 1, n]) / h[k] ** 2 

314 

315 def _total_time_cost(self, h, weight=1.0): 

316 """ 

317 Helper function that sets a quadratic cost on the total trajectory time. 

318 

319 Parameters 

320 ---------- 

321 h : pydrake.autodiffutils.AutoDiffXd 

322 The time segment durations. 

323 weight : float, optional 

324 A weight to scale the cost function. 

325 

326 Return 

327 ------ 

328 pydrake.autodiffutils.AutoDiffXd 

329 Quadratic cost on the total trajectory time. 

330 """ 

331 total_time = 0.0 

332 for step in h: 

333 total_time += step 

334 return weight * total_time**2 

335 

336 def _collision_constraint( 

337 self, 

338 vars, 

339 num_waypoints, 

340 num_dofs, 

341 samples_per_segment, 

342 collision_pairs, 

343 min_allowable_dist, 

344 influence_dist, 

345 ): 

346 """ 

347 Helper function to evaluate collision constraint and its gradients. 

348 

349 These calculations are based off the following resources: 

350 * https://typeset.io/pdf/a-collision-free-mpc-for-whole-body-dynamic-locomotion-and-1l6itpfk.pdf 

351 * https://laas.hal.science/hal-04425002 

352 

353 Parameters 

354 ---------- 

355 vars : pydrake.autodiffutils.AutoDiffXd 

356 A flattened list of joint positions, joint velocities, collocation positions, 

357 and waypoint times. Needs to be in one variable to add this function as a constraint. 

358 num_waypoints : int 

359 The number of waypoints. 

360 num_dofs : int 

361 The number of degrees of freedom. 

362 samples_per_segment : int 

363 The number of samples per waypoint at which to evaluate collisions. 

364 collision_pairs : list[int] 

365 A list containing the indices of collision model pairs to check. 

366 min_allowable_dist : float 

367 The minimum allowable collision distance, in meters, for the constraint to be satisfied. 

368 Positive is out of collision and negative is in collision. 

369 influence_dist : float 

370 The distance over which the collision distance has an effect on the constraint gradient. 

371 This value should be larger than `min_allowable_dist`. 

372 

373 Return 

374 ------ 

375 pydrake.autodiffutils.AutoDiffXd 

376 The constraint value representing the minimum collision distance for every collision object and waypoint, 

377 as well as its corresponding gradient. 

378 """ 

379 all_vars = ExtractValue(vars) 

380 all_gradients = ExtractGradient(vars) 

381 x_all = all_vars[: num_waypoints * num_dofs].reshape((num_waypoints, num_dofs)) 

382 x_d_all = all_vars[ 

383 num_waypoints * num_dofs : 2 * num_waypoints * num_dofs 

384 ].reshape((num_waypoints, num_dofs)) 

385 xc_d_all = all_vars[ 

386 2 

387 * num_waypoints 

388 * num_dofs : (2 * num_waypoints + num_waypoints - 1) 

389 * num_dofs 

390 ].reshape((num_waypoints - 1, num_dofs)) 

391 h_all = all_vars[-(num_waypoints - 1) :] 

392 

393 num_collision_pairs = len(collision_pairs) 

394 total_num_points = (num_waypoints - 1) * samples_per_segment 

395 

396 min_distance_smoothed_squared = np.zeros(total_num_points * num_collision_pairs) 

397 

398 gradient = np.zeros( 

399 (total_num_points * num_collision_pairs, all_gradients.shape[0]) 

400 ) 

401 point_idx = 0 

402 for k in range(num_waypoints - 1): 

403 for step in np.linspace(0.0, 1.0, samples_per_segment): 

404 q = np.array( 

405 [ 

406 self._eval_position(x_all, x_d_all, xc_d_all, h_all, k, n, step) 

407 for n in range(num_dofs) 

408 ] 

409 ) 

410 

411 # Compute collision and distance checks 

412 pinocchio.framesForwardKinematics(self.model, self.data, q) 

413 pinocchio.computeDistances( 

414 self.model, self.data, self.collision_model, self.collision_data, q 

415 ) 

416 

417 # Get the minimum distance in the allowable range 

418 for link_idx, pair in enumerate(collision_pairs): 

419 dist = self.collision_data.distanceResults[pair].min_distance 

420 

421 # Find the collision Jacobian for the closest point pair, and calculate gradients. 

422 if dist <= influence_dist: 

423 distance_vec, Jcoll1, Jcoll2 = ( 

424 calculate_collision_vector_and_jacobians( 

425 self.model, 

426 self.collision_model, 

427 self.data, 

428 self.collision_data, 

429 pair, 

430 q, 

431 ) 

432 ) 

433 grad = distance_vec.T @ (Jcoll2 - Jcoll1) 

434 else: 

435 grad = np.zeros(num_dofs) 

436 

437 # Minimum distances are smoothed using a squared hinge loss to have better gradients. 

438 min_distance_smoothed = np.clip( 

439 1.0 

440 + (dist - min_allowable_dist) 

441 / (min_allowable_dist - influence_dist), 

442 0.0, 

443 np.inf, 

444 ) 

445 grad_idx = point_idx * num_collision_pairs + link_idx 

446 min_distance_smoothed_squared[grad_idx] = min_distance_smoothed**2 

447 grad = ( 

448 2.0 

449 * grad 

450 * min_distance_smoothed 

451 / (min_allowable_dist - influence_dist) 

452 ) 

453 

454 # Assign the gradients for x, x_d, xc_d, and h at this index. 

455 start = k * num_dofs 

456 end = start + num_dofs 

457 gradient[grad_idx, start:end] += grad 

458 

459 start = (num_waypoints * num_dofs) + k * num_dofs 

460 end = start + num_dofs 

461 gradient[grad_idx, start:end] += grad 

462 

463 start = (num_waypoints * 2 * num_dofs) + k * num_dofs 

464 end = start + num_dofs 

465 gradient[grad_idx, start:end] += grad 

466 

467 h_idx = ((num_waypoints * 2 + num_waypoints - 1) * num_dofs) + k 

468 gradient[grad_idx, h_idx] += np.sum(grad) 

469 

470 point_idx += 1 

471 

472 return InitializeAutoDiff( 

473 min_distance_smoothed_squared, gradient @ all_gradients 

474 ) 

475 

476 def plan(self, q_path, init_path=None): 

477 """ 

478 Plans a trajectory from a start to a goal configuration, or along an entire trajectory. 

479 

480 If the input list has 2 elements, then this is assumed to be the start and goal configurations. 

481 The intermediate waypoints will be determined automatically. 

482 

483 If the input list has more than 2 elements, then these are the actual waypoints that must be achieved. 

484 

485 Parameters 

486 ---------- 

487 q_path : list[array-like] 

488 A list of joint configurations describing the desired motion. 

489 init_path : list[array-like], optional 

490 If set, defines the initial guess for the path waypoints. 

491 

492 Return 

493 ------ 

494 Optional[pyroboplan.trajectory.polynomial.CubicPolynomialTrajectory] 

495 The resulting trajectory, or None if optimization failed 

496 """ 

497 if len(q_path) == 0: 

498 warnings.warn("Cannot optimize over an empty path.") 

499 return None 

500 num_waypoints = self.options.num_waypoints 

501 num_dofs = len(q_path[0]) 

502 

503 if len(q_path) == num_waypoints: 

504 fully_specified_path = True 

505 elif len(q_path) == 2: 

506 fully_specified_path = False 

507 else: 

508 raise ValueError("Path must either be length 2 or equal to num_waypoints.") 

509 

510 # Preprocess the kinematic limits. 

511 min_vel = self._process_limits(self.options.min_vel, num_dofs, "min_vel") 

512 max_vel = self._process_limits(self.options.max_vel, num_dofs, "max_vel") 

513 min_accel = self._process_limits(self.options.min_accel, num_dofs, "min_accel") 

514 max_accel = self._process_limits(self.options.max_accel, num_dofs, "max_accel") 

515 min_jerk = self._process_limits(self.options.min_jerk, num_dofs, "min_jerk") 

516 max_jerk = self._process_limits(self.options.max_jerk, num_dofs, "max_jerk") 

517 

518 # Initialize the basic program and its variables. 

519 prog = MathematicalProgram() 

520 x = prog.NewContinuousVariables(num_waypoints, num_dofs) 

521 x_d = prog.NewContinuousVariables(num_waypoints, num_dofs) 

522 xc = prog.NewContinuousVariables(num_waypoints - 1, num_dofs) 

523 xc_d = prog.NewContinuousVariables(num_waypoints - 1, num_dofs) 

524 h = prog.NewContinuousVariables(num_waypoints - 1) 

525 

526 # Initial, final, and intermediate waypoint conditions. 

527 if fully_specified_path: 

528 for idx in range(num_waypoints): 

529 prog.AddBoundingBoxConstraint(q_path[idx], q_path[idx], x[idx, :]) 

530 else: 

531 prog.AddBoundingBoxConstraint(q_path[0], q_path[0], x[0, :]) 

532 prog.AddBoundingBoxConstraint(q_path[-1], q_path[-1], x[-1, :]) 

533 # Initial and final velocities should always be zero. 

534 prog.AddBoundingBoxConstraint(0.0, 0.0, x_d[0, :]) 

535 prog.AddBoundingBoxConstraint(0.0, 0.0, x_d[num_waypoints - 1, :]) 

536 

537 for n in range(num_dofs): 

538 # Collocation point constraints. 

539 # Specifically, this constrains the position and velocities of the collocation points to be 

540 # expressed in terms of the waypoint positions and velocities, assuming cubic splines. 

541 # This is described in the "Direct Collocation" section here: 

542 # https://underactuated.mit.edu/trajopt.html 

543 for k in range(num_waypoints - 1): 

544 prog.AddQuadraticConstraint( 

545 0.5 * (x[k, n] + x[k + 1, n]) 

546 + (h[k] / 8.0) * (x_d[k, n] - x_d[k + 1, n]) 

547 - xc[k, n], 

548 0.0, 

549 0.0, 

550 QuadraticConstraint.HessianType.kIndefinite, 

551 ) 

552 prog.AddConstraint( 

553 xc_d[k, n] 

554 == -(1.5 / h[k]) * (x[k, n] - x[k + 1, n]) 

555 - 0.25 * (x_d[k, n] + x_d[k + 1, n]) 

556 ) 

557 

558 # Sample points along each segment to evaluate the position and velocity constraints, since 

559 # they are cubic and quadratic (respectively) and could overshoot the waypoints and collocation points. 

560 for step in np.linspace(0.0, 1.0, self.options.samples_per_segment): 

561 pos = self._eval_position(x, x_d, xc_d, h, k, n, step) 

562 prog.AddConstraint(pos <= self.model.upperPositionLimit[n]) 

563 prog.AddConstraint(pos >= self.model.lowerPositionLimit[n]) 

564 

565 vel = self._eval_velocity(x_d, xc_d, h, k, n, step) 

566 prog.AddConstraint(vel >= min_vel[n]) 

567 prog.AddConstraint(vel <= max_vel[n]) 

568 

569 # The acceleration and jerk are linear and piecewise constant (respectively), 

570 # so these can be constrained simply by checking each endpoint. 

571 for step in [0.0, 1.0]: 

572 accel = self._eval_acceleration(x_d, xc_d, h, k, n, step) 

573 prog.AddConstraint(accel >= min_accel[n]) 

574 prog.AddConstraint(accel <= max_accel[n]) 

575 

576 jerk = self._eval_jerk(x_d, xc_d, h, k, n, step) 

577 prog.AddConstraint(jerk >= min_jerk[n]) 

578 prog.AddConstraint(jerk <= max_jerk[n]) 

579 

580 # Enforce acceleration continuity between segments. 

581 # That is, the final acceleration of the kth waypoint must equal the initial acceleration of the (k+1)th waypoint. 

582 for k in range(num_waypoints - 2): 

583 prog.AddConstraint( 

584 self._eval_acceleration(x_d, xc_d, h, k, n, 1.0) 

585 == self._eval_acceleration(x_d, xc_d, h, k + 1, n, 0.0) 

586 ) 

587 

588 # Cost and bounds on trajectory times. 

589 prog.AddBoundingBoxConstraint( 

590 self.options.min_segment_time, self.options.max_segment_time, h 

591 ) 

592 prog.AddCost(self._total_time_cost, h) 

593 

594 # Collision checking by sampling points along the entire trajectory. 

595 link_list = self.options.collision_link_list 

596 if self.options.check_collisions and len(link_list) > 0: 

597 all_pairs = get_collision_pair_indices_from_bodies( 

598 self.model, self.collision_model, link_list 

599 ) 

600 num_pairs = len(all_pairs) 

601 

602 total_points = (num_waypoints - 1) * self.options.samples_per_segment 

603 # The minimum distance values are smoothed in the constraint such that 

604 # a value of 1.0 or lower satisfies the constraint. 

605 min_dist_val = -np.inf * np.ones(total_points * num_pairs) 

606 max_dist_val = 1.0 * np.ones(total_points * num_pairs) 

607 collision_expr = partial( 

608 self._collision_constraint, 

609 num_waypoints=num_waypoints, 

610 num_dofs=num_dofs, 

611 samples_per_segment=self.options.samples_per_segment, 

612 collision_pairs=all_pairs, 

613 min_allowable_dist=self.options.min_collision_dist, 

614 influence_dist=self.options.collision_influence_dist, 

615 ) 

616 vars = np.concat([x.flatten(), x_d.flatten(), xc_d.flatten(), h.flatten()]) 

617 prog.AddConstraint(collision_expr, min_dist_val, max_dist_val, vars) 

618 

619 # Optionally add a cost for collision avoidance. 

620 # Since the smoothed minimum distance is 1.0 at the constraint limit and 

621 # increases further into collision, we want to minimize this value. 

622 coll_cost_weight = self.options.collision_avoidance_cost_weight 

623 if coll_cost_weight > 0.0: 623 ↛ 624line 623 didn't jump to line 624 because the condition on line 623 was never true

624 prog.AddCost( 

625 lambda var: -coll_cost_weight * max(collision_expr(var))[0], 

626 vars, 

627 ) 

628 

629 # Set initial conditions to help search. 

630 if init_path or fully_specified_path: 

631 # Set initial guess assuming collocation points are exactly between the waypoints. 

632 if init_path: 

633 q_path = init_path 

634 prog.SetInitialGuess(x, np.array(q_path)) 

635 init_collocation_points = [] 

636 for k in range(num_waypoints - 1): 

637 init_collocation_points.append(0.5 * (q_path[k] + q_path[k + 1])) 

638 prog.SetInitialGuess(xc, np.array(init_collocation_points)) 

639 else: 

640 # Set initial guess assuming linear trajectory from start to end. 

641 init_points = np.linspace(q_path[0], q_path[-1], 2 * num_waypoints - 1) 

642 prog.SetInitialGuess(x, init_points[::2]) 

643 prog.SetInitialGuess(xc, init_points[1::2]) 

644 

645 h_init = 0.5 * (self.options.min_segment_time + self.options.max_segment_time) 

646 prog.SetInitialGuess(h, h_init * np.ones(num_waypoints - 1)) 

647 prog.SetInitialGuess(x_d, np.zeros((num_waypoints, num_dofs))) 

648 prog.SetInitialGuess(xc_d, np.zeros((num_waypoints - 1, num_dofs))) 

649 

650 # Solve the program. 

651 solver_options = SolverOptions() 

652 solver_options.SetOption( 

653 SnoptSolver.id(), "Time limit", self.options.max_planning_time 

654 ) 

655 solver_options.SetOption(SnoptSolver.id(), "Timing level", 3) 

656 result = Solve(prog, solver_options=solver_options) 

657 if not result.is_success(): 

658 print("Trajectory optimization failed.") 

659 return None 

660 

661 # Unpack the values. 

662 h_opt = result.GetSolution(h) 

663 x_opt = result.GetSolution(x) 

664 x_d_opt = result.GetSolution(x_d) 

665 

666 # Generate the cubic trajectory and return it. 

667 t_opt = [0] + list(np.cumsum(h_opt)) 

668 self.latest_trajectory = CubicPolynomialTrajectory( 

669 np.array(t_opt), 

670 np.array(x_opt.T), 

671 np.array(x_d_opt.T), 

672 ) 

673 return self.latest_trajectory