API  4.5
For C++ developers
examplePredictAndTrack.py

This example predicts the trajectory to swing a double pendulum from horizontal to vertical, and then tracks the resulting trajectory.

1 # -------------------------------------------------------------------------- #
2 # OpenSim Moco: examplePredictAndTrack.py #
3 # -------------------------------------------------------------------------- #
4 # Copyright (c) 2018 Stanford University and the Authors #
5 # #
6 # Author(s): Christopher Dembia #
7 # #
8 # Licensed under the Apache License, Version 2.0 (the "License"); you may #
9 # not use this file except in compliance with the License. You may obtain a #
10 # copy of the License at http://www.apache.org/licenses/LICENSE-2.0 #
11 # #
12 # Unless required by applicable law or agreed to in writing, software #
13 # distributed under the License is distributed on an "AS IS" BASIS, #
14 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. #
15 # See the License for the specific language governing permissions and #
16 # limitations under the License. #
17 # -------------------------------------------------------------------------- #
18 
19 import os
20 import math
21 import opensim as osim
22 
23 """
24 This file performs the following problems using a
25 double pendulum model:
26 
27  1. predict an optimal trajectory (and controls),
28  2. track the states from the optimal trajectory, and
29  3. track the marker trajectories from the optimal trajectory.
30 
31 """
32 
33 visualize = True
34 # The following environment variable is set during automated testing.
35 if os.getenv('OPENSIM_USE_VISUALIZER') == '0':
36  visualize = False
37 
38 # Create a model of a double pendulum.
39 # ------------------------------------
40 def createDoublePendulumModel():
41  model = osim.Model()
42  model.setName("double_pendulum")
43 
44  # Create two links, each with a mass of 1 kg, center of mass at the body's
45  # origin, and moments and products of inertia of zero.
46  b0 = osim.Body("b0", 1, osim.Vec3(0), osim.Inertia(1))
47  model.addBody(b0)
48  b1 = osim.Body("b1", 1, osim.Vec3(0), osim.Inertia(1))
49  model.addBody(b1)
50 
51  # Add markers to body origin locations.
52  m0 = osim.Marker("m0", b0, osim.Vec3(0))
53  m1 = osim.Marker("m1", b1, osim.Vec3(0))
54  model.addMarker(m0)
55  model.addMarker(m1)
56 
57  # Connect the bodies with pin joints. Assume each body is 1 m long.
58  j0 = osim.PinJoint("j0", model.getGround(), osim.Vec3(0), osim.Vec3(0),
59  b0, osim.Vec3(-1, 0, 0), osim.Vec3(0))
60  q0 = j0.updCoordinate()
61  q0.setName("q0")
62  j1 = osim.PinJoint("j1",
63  b0, osim.Vec3(0), osim.Vec3(0), b1, osim.Vec3(-1, 0, 0), osim.Vec3(0))
64  q1 = j1.updCoordinate()
65  q1.setName("q1")
66  model.addJoint(j0)
67  model.addJoint(j1)
68 
69  tau0 = osim.CoordinateActuator()
70  tau0.setCoordinate(j0.updCoordinate())
71  tau0.setName("tau0")
72  tau0.setOptimalForce(1)
73  model.addComponent(tau0)
74 
75  tau1 = osim.CoordinateActuator()
76  tau1.setCoordinate(j1.updCoordinate())
77  tau1.setName("tau1")
78  tau1.setOptimalForce(1)
79  model.addComponent(tau1)
80 
81  # Add display geometry.
82  bodyGeometry = osim.Ellipsoid(0.5, 0.1, 0.1)
83  transform = osim.Transform(osim.Vec3(-0.5, 0, 0))
84  b0Center = osim.PhysicalOffsetFrame("b0_center", b0, transform)
85  b0.addComponent(b0Center)
86  b0Center.attachGeometry(bodyGeometry.clone())
87  b1Center = osim.PhysicalOffsetFrame("b1_center", b1, transform)
88  b1.addComponent(b1Center)
89  b1Center.attachGeometry(bodyGeometry.clone())
90 
91  model.finalizeConnections()
92  model.printToXML("double_pendulum.osim")
93  return model
94 
95 
96 def solvePrediction():
97  # Predict the optimal trajectory for a minimum time swing-up.
98  # In the diagram below, + represents the origin, and ---o represents a link
99  # in the double pendulum.
100  #
101  # o
102  # |
103  # o
104  # |
105  # +---o---o +
106  #
107  # iniital pose final pose
108  #
109  study = osim.MocoStudy()
110  study.setName("double_pendulum_predict")
111 
112  problem = study.updProblem()
113 
114  # Model (dynamics).
115  problem.setModel(createDoublePendulumModel())
116 
117  # Bounds.
118  problem.setTimeBounds(0, [0, 5])
119  # Arguments are name, [lower bound, upper bound],
120  # initial [lower bound, upper bound],
121  # final [lower bound, upper bound].
122  problem.setStateInfo("/jointset/j0/q0/value", [-10, 10], 0)
123  problem.setStateInfo("/jointset/j0/q0/speed", [-50, 50], 0, 0)
124  problem.setStateInfo("/jointset/j1/q1/value", [-10, 10], 0)
125  problem.setStateInfo("/jointset/j1/q1/speed", [-50, 50], 0, 0)
126  problem.setControlInfo("/tau0", [-100, 100])
127  problem.setControlInfo("/tau1", [-100, 100])
128 
129  # Cost: minimize final time and error from desired
130  # end effector position.
131  ftCost = osim.MocoFinalTimeGoal()
132  ftCost.setWeight(0.001)
133  problem.addGoal(ftCost)
134 
135  finalCost = osim.MocoMarkerFinalGoal()
136  finalCost.setName("final")
137  finalCost.setWeight(1000.0)
138  finalCost.setPointName("/markerset/m1")
139  finalCost.setReferenceLocation(osim.Vec3(0, 2, 0))
140  problem.addGoal(finalCost)
141 
142 
143  # Configure the solver.
144  solver = study.initTropterSolver()
145  solver.set_num_mesh_intervals(100)
146  solver.set_verbosity(2)
147  solver.set_optim_solver("ipopt")
148 
149  guess = solver.createGuess()
150  guess.setNumTimes(2)
151  guess.setTime([0, 1])
152  guess.setState("/jointset/j0/q0/value", [0, -math.pi])
153  guess.setState("/jointset/j1/q1/value", [0, 2*math.pi])
154  guess.setState("/jointset/j0/q0/speed", [0, 0])
155  guess.setState("/jointset/j1/q1/speed", [0, 0])
156  guess.setControl("/tau0", [0, 0])
157  guess.setControl("/tau1", [0, 0])
158  guess.resampleWithNumTimes(10)
159  solver.setGuess(guess)
160 
161  # Save the problem to a setup file for reference.
162  study.printToXML("examplePredictAndTrack_predict.omoco")
163 
164  # Solve the problem.
165  solution = study.solve()
166  solution.write("examplePredictAndTrack_predict_solution.sto")
167 
168  if visualize:
169  study.visualize(solution)
170  return solution
171 
172 
173 def computeMarkersReference(predictedSolution):
174  model = createDoublePendulumModel()
175  model.initSystem()
176  states = predictedSolution.exportToStatesTable()
177 
178  statesTraj = osim.StatesTrajectory.createFromStatesTable(model, states)
179 
180  markerTrajectories = osim.TimeSeriesTableVec3()
181  markerTrajectories.setColumnLabels(["/markerset/m0", "/markerset/m1"])
182 
183  for state in statesTraj:
184  model.realizePosition(state)
185  m0 = model.getComponent("markerset/m0")
186  m1 = model.getComponent("markerset/m1")
187  markerTrajectories.appendRow(state.getTime(),
188  osim.RowVectorVec3([m0.getLocationInGround(state),
189  m1.getLocationInGround(state)]))
190 
191  # Assign a weight to each marker.
192  markerWeights = osim.SetMarkerWeights()
193  markerWeights.cloneAndAppend(osim.MarkerWeight("/markerset/m0", 1))
194  markerWeights.cloneAndAppend(osim.MarkerWeight("/markerset/m1", 5))
195 
196  return osim.MarkersReference(markerTrajectories, markerWeights)
197 
198 
199 def solveStateTracking(stateRef):
200  # Predict the optimal trajectory for a minimum time swing-up.
201  study = osim.MocoStudy()
202  study.setName("double_pendulum_track")
203 
204  problem = study.updProblem()
205 
206  # Model (dynamics).
207  problem.setModel(createDoublePendulumModel())
208 
209  # Bounds.
210  # Arguments are name, [lower bound, upper bound],
211  # initial [lower bound, upper bound],
212  # final [lower bound, upper bound].
213  finalTime = stateRef.getIndependentColumn()[-1]
214  problem.setTimeBounds(0, finalTime)
215  problem.setStateInfo("/jointset/j0/q0/value", [-10, 10], 0)
216  problem.setStateInfo("/jointset/j0/q0/speed", [-50, 50], 0)
217  problem.setStateInfo("/jointset/j1/q1/value", [-10, 10], 0)
218  problem.setStateInfo("/jointset/j1/q1/speed", [-50, 50], 0)
219  problem.setControlInfo("/tau0", [-150, 150])
220  problem.setControlInfo("/tau1", [-150, 150])
221 
222  # Cost: track provided state data.
223  stateTracking = osim.MocoStateTrackingGoal()
224  stateTracking.setReference(osim.TableProcessor(stateRef))
225  problem.addGoal(stateTracking)
226 
227  effort = osim.MocoControlGoal()
228  effort.setName("effort")
229  effort.setWeight(0.001)
230  # TODO problem.addGoal(effort)
231 
232  # Configure the solver.
233  solver = study.initTropterSolver()
234  solver.set_num_mesh_intervals(50)
235  solver.set_verbosity(2)
236  solver.set_optim_solver("ipopt")
237  solver.set_optim_jacobian_approximation("exact")
238  solver.set_optim_hessian_approximation("exact")
239  solver.set_exact_hessian_block_sparsity_mode("dense")
240 
241  # Save the problem to a setup file for reference.
242  study.printToXML("examplePredictAndTrack_track_states.omoco")
243 
244  # Solve the problem.
245  solution = study.solve()
246 
247  solution.write("examplePredictAndTrack_track_states_solution.sto")
248 
249  if visualize:
250  study.visualize(solution)
251  return solution
252 
253 
254 def solveMarkerTracking(markersRef, guess):
255  # Predict the optimal trajectory for a minimum time swing-up.
256  study = osim.MocoStudy()
257  study.setName("double_pendulum_track")
258 
259  problem = study.updProblem()
260 
261  # Model (dynamics).
262  problem.setModel(createDoublePendulumModel())
263 
264  # Bounds.
265  # Arguments are name, [lower bound, upper bound],
266  # initial [lower bound, upper bound],
267  # final [lower bound, upper bound].
268  finalTime = markersRef.getMarkerTable().getIndependentColumn()[-1]
269  problem.setTimeBounds(0, finalTime)
270  problem.setStateInfo("/jointset/j0/q0/value", [-10, 10], 0)
271  problem.setStateInfo("/jointset/j0/q0/speed", [-50, 50], 0)
272  problem.setStateInfo("/jointset/j1/q1/value", [-10, 10], 0)
273  problem.setStateInfo("/jointset/j1/q1/speed", [-50, 50], 0)
274  problem.setControlInfo("/tau0", [-100, 100])
275  problem.setControlInfo("/tau1", [-100, 100])
276 
277  # Cost: track provided marker data.
278  markerTracking = osim.MocoMarkerTrackingGoal()
279  markerTracking.setMarkersReference(markersRef)
280  problem.addGoal(markerTracking)
281 
282  effort = osim.MocoControlGoal()
283  effort.setName("effort")
284  effort.setWeight(0.0001)
285  # problem.addGoal(effort)
286 
287  # Configure the solver.
288  solver = study.initTropterSolver()
289  solver.set_num_mesh_intervals(50)
290  solver.set_verbosity(2)
291  solver.set_optim_solver("ipopt")
292  solver.set_optim_jacobian_approximation("exact")
293  solver.set_optim_hessian_approximation("exact")
294  solver.set_exact_hessian_block_sparsity_mode("dense")
295 
296  solver.setGuess(guess)
297 
298  # Save the problem to a setup file for reference.
299  study.printToXML("examplePredictAndTrack_track_markers.omoco")
300 
301  # Solve the problem.
302  solution = study.solve()
303 
304  solution.write("examplePredictAndTrack_track_markers_solution.sto")
305 
306  if visualize:
307  study.visualize(solution)
308  return solution
309 
310 
311 
312 
313 optimalTrajectory = solvePrediction()
314 
315 markersRef = computeMarkersReference(optimalTrajectory)
316 
317 trackedSolution = solveStateTracking(optimalTrajectory.exportToStatesTable())
318 
319 trackedSolution2 = solveMarkerTracking(markersRef, trackedSolution)
320 
321 
322 
323 
324 
325 
326 
327 
328 
329 
330 
331 
332 
333 
334 
335 
336