openAIgymTriplePendulum.py
You can view and download this file on Github: openAIgymTriplePendulum.py
1#+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
2# This is an EXUDYN example
3#
4# Details: This file shows integration with OpenAI gym by testing a triple pendulum example
5#
6# Author: Johannes Gerstmayr
7# Date: 2022-05-18
8#
9# Copyright:This file is part of Exudyn. Exudyn is free software. You can redistribute it and/or modify it under the terms of the Exudyn license. See 'LICENSE.txt' for more details.
10#
11#+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
12
13
14import exudyn as exu
15from exudyn.utilities import * #includes itemInterface and rigidBodyUtilities
16import numpy as np
17from math import sin, cos
18import exudyn.graphics as graphics
19from exudyn.artificialIntelligence import *
20import math
21import os
22
23class InvertedTriplePendulumEnv(OpenAIGymInterfaceEnv):
24
25 #**classFunction: OVERRIDE this function to create multibody system mbs and setup simulationSettings; call Assemble() at the end!
26 # you may also change SC.visualizationSettings() individually; kwargs may be used for special setup
27 def CreateMBS(self, SC, mbs, simulationSettings, **kwargs):
28
29 #%%++++++++++++++++++++++++++++++++++++++++++++++
30 #this model uses kwargs: thresholdFactor
31 thresholdFactor = 3
32 if 'thresholdFactor' in kwargs:
33 thresholdFactor = kwargs['thresholdFactor']
34
35 gravity = 9.81
36 self.length = 1.
37 width = 0.1*self.length
38 masscart = 1.
39 massarm = 0.1
40 total_mass = massarm + masscart
41 armInertia = self.length**2*0.5*massarm
42 self.force_mag = 10.0*2 #must be larger for triple pendulum to be more reactive ...
43 self.stepUpdateTime = 0.02 # seconds between state updates
44
45 background = graphics.CheckerBoard(point= [0,0.5*self.length,-0.5*width],
46 normal= [0,0,1], size=10, size2=6, nTiles=20, nTiles2=12)
47
48 oGround=self.mbs.AddObject(ObjectGround(referencePosition= [0,0,0], #x-pos,y-pos,angle
49 visualization=VObjectGround(graphicsData= [background])))
50 nGround=self.mbs.AddNode(NodePointGround())
51
52 gCart = graphics.Brick(size=[0.5*self.length, width, width],
53 color=graphics.color.dodgerblue)
54 self.nCart = self.mbs.AddNode(Rigid2D(referenceCoordinates=[0,0,0]));
55 oCart = self.mbs.AddObject(RigidBody2D(mass=masscart,
56 inertia=0.1*masscart, #not needed
57 nodeNumber=self.nCart,
58 visualization=VObjectRigidBody2D(graphicsData= [gCart])))
59 mCartCOM = self.mbs.AddMarker(MarkerNodePosition(nodeNumber=self.nCart))
60
61 gArm1 = graphics.Brick(size=[width, self.length, width], color=graphics.color.red)
62 gArm1joint = graphics.Cylinder(pAxis=[0,-0.5*self.length,-0.6*width], vAxis=[0,0,1.2*width],
63 radius=0.0625*self.length, color=graphics.color.darkgrey)
64 self.nArm1 = self.mbs.AddNode(Rigid2D(referenceCoordinates=[0,0.5*self.length,0]));
65 oArm1 = self.mbs.AddObject(RigidBody2D(mass=massarm,
66 inertia=armInertia, #not included in original paper
67 nodeNumber=self.nArm1,
68 visualization=VObjectRigidBody2D(graphicsData= [gArm1, gArm1joint])))
69
70 mArm1COM = self.mbs.AddMarker(MarkerNodePosition(nodeNumber=self.nArm1))
71 mArm1JointA = self.mbs.AddMarker(MarkerBodyPosition(bodyNumber=oArm1, localPosition=[0,-0.5*self.length,0]))
72 mArm1JointB = self.mbs.AddMarker(MarkerBodyPosition(bodyNumber=oArm1, localPosition=[0, 0.5*self.length,0]))
73
74 gArm2 = graphics.Brick(size=[width, self.length, width], color=graphics.color.red)
75 self.nArm2 = self.mbs.AddNode(Rigid2D(referenceCoordinates=[0,1.5*self.length,0]));
76 oArm2 = self.mbs.AddObject(RigidBody2D(mass=massarm,
77 inertia=armInertia, #not included in original paper
78 nodeNumber=self.nArm2,
79 visualization=VObjectRigidBody2D(graphicsData= [gArm2, gArm1joint])))
80
81 mArm2COM = self.mbs.AddMarker(MarkerNodePosition(nodeNumber=self.nArm2))
82 mArm2Joint = self.mbs.AddMarker(MarkerBodyPosition(bodyNumber=oArm2, localPosition=[0,-0.5*self.length,0]))
83 mArm2JointB = self.mbs.AddMarker(MarkerBodyPosition(bodyNumber=oArm2, localPosition=[0, 0.5*self.length,0]))
84
85 gArm3 = graphics.Brick(size=[width, self.length, width], color=graphics.color.red)
86 self.nArm3 = self.mbs.AddNode(Rigid2D(referenceCoordinates=[0,2.5*self.length,0]));
87 oArm3 = self.mbs.AddObject(RigidBody2D(mass=massarm,
88 inertia=armInertia, #not included in original paper
89 nodeNumber=self.nArm3,
90 visualization=VObjectRigidBody2D(graphicsData= [gArm3, gArm1joint])))
91
92 mArm3COM = self.mbs.AddMarker(MarkerNodePosition(nodeNumber=self.nArm3))
93 mArm3Joint = self.mbs.AddMarker(MarkerBodyPosition(bodyNumber=oArm3, localPosition=[0,-0.5*self.length,0]))
94
95 mCartCoordX = self.mbs.AddMarker(MarkerNodeCoordinate(nodeNumber=self.nCart, coordinate=0))
96 mCartCoordY = self.mbs.AddMarker(MarkerNodeCoordinate(nodeNumber=self.nCart, coordinate=1))
97 mGroundNode = self.mbs.AddMarker(MarkerNodeCoordinate(nodeNumber=nGround, coordinate=0))
98
99 #gravity
100 self.mbs.AddLoad(Force(markerNumber=mCartCOM, loadVector=[0,-masscart*gravity,0]))
101 self.mbs.AddLoad(Force(markerNumber=mArm1COM, loadVector=[0,-massarm*gravity,0]))
102 self.mbs.AddLoad(Force(markerNumber=mArm2COM, loadVector=[0,-massarm*gravity,0]))
103 self.mbs.AddLoad(Force(markerNumber=mArm3COM, loadVector=[0,-massarm*gravity,0]))
104
105 #control force
106 self.lControl = self.mbs.AddLoad(LoadCoordinate(markerNumber=mCartCoordX, load=1.))
107
108 #joints and constraints:
109 self.mbs.AddObject(RevoluteJoint2D(markerNumbers=[mCartCOM, mArm1JointA]))
110 self.mbs.AddObject(RevoluteJoint2D(markerNumbers=[mArm1JointB, mArm2Joint]))
111 self.mbs.AddObject(RevoluteJoint2D(markerNumbers=[mArm2JointB, mArm3Joint]))
112
113 self.mbs.AddObject(CoordinateConstraint(markerNumbers=[mCartCoordY, mGroundNode]))
114
115
116
117
118 #%%++++++++++++++++++++++++
119 self.mbs.Assemble() #computes initial vector
120
121 self.simulationSettings.timeIntegration.numberOfSteps = 1
122 self.simulationSettings.timeIntegration.endTime = 0 #will be overwritten in step
123 self.simulationSettings.timeIntegration.verboseMode = 0
124 self.simulationSettings.solution.file.write = False
125 #self.simulationSettings.timeIntegration.realtime.active = True
126
127 self.simulationSettings.timeIntegration.newton.useModifiedNewton = True
128
129 self.SC.visualizationSettings.view0.scene.drawWorldBasis=True
130 self.SC.visualizationSettings.general.graphicsUpdateInterval = 0.01 #50Hz
131 self.SC.visualizationSettings.openGL.multiSampling=4
132
133 #self.simulationSettings.solution.file.information = "Open AI gym"
134
135 #+++++++++++++++++++++++++++++++++++++++++++++++++++++
136 # Angle at which to fail the episode
137 # these parameters are used in subfunctions
138 self.theta_threshold_radians = thresholdFactor* 12 * 2 * math.pi / 360
139 self.x_threshold = thresholdFactor*2.4
140
141 #must return state size
142 stateSize = 8 #the number of states (position/velocity that are used by learning algorithm)
143 return stateSize
144
145 #**classFunction: OVERRIDE this function to set up self.action_space and self.observation_space
146 def SetupSpaces(self):
147
148 high = np.array(
149 [
150 self.x_threshold * 2,
151 np.finfo(np.float32).max,
152 self.theta_threshold_radians * 2,
153 np.finfo(np.float32).max,
154 self.theta_threshold_radians * 2,
155 np.finfo(np.float32).max,
156 self.theta_threshold_radians * 2,
157 np.finfo(np.float32).max,
158 ],
159 dtype=np.float32,
160 )
161
162 #+++++++++++++++++++++++++++++++++++++++++++++++++++++
163 #see https://github.com/openai/gym/blob/64b4b31d8245f6972b3d37270faf69b74908a67d/gym/core.py#L16
164 #for Env:
165 self.action_space = spaces.Discrete(2)
166 self.observation_space = spaces.Box(-high, high, dtype=np.float32)
167 #+++++++++++++++++++++++++++++++++++++++++++++++++++++
168
169
170 #**classFunction: OVERRIDE this function to map the action given by learning algorithm to the multibody system, e.g. as a load parameter
171 def MapAction2MBS(self, action):
172 force = self.force_mag if action == 1 else -self.force_mag
173 self.mbs.SetLoadParameter(self.lControl, 'load', force)
174
175 #**classFunction: OVERRIDE this function to collect output of simulation and map to self.state tuple
176 #**output: return bool done which contains information if system state is outside valid range
177 def Output2StateAndDone(self):
178
179 #+++++++++++++++++++++++++
180 #compute some output:
181 cartPosX = self.mbs.GetNodeOutput(self.nCart, variableType=exu.OutputVariableType.Coordinates)[0]
182 arm1Angle = self.mbs.GetNodeOutput(self.nArm1, variableType=exu.OutputVariableType.Coordinates)[2]
183 arm2Angle = self.mbs.GetNodeOutput(self.nArm2, variableType=exu.OutputVariableType.Coordinates)[2]
184 arm3Angle = self.mbs.GetNodeOutput(self.nArm3, variableType=exu.OutputVariableType.Coordinates)[2]
185 cartPosX_t = self.mbs.GetNodeOutput(self.nCart, variableType=exu.OutputVariableType.Coordinates_t)[0]
186 arm1Angle_t = self.mbs.GetNodeOutput(self.nArm1, variableType=exu.OutputVariableType.Coordinates_t)[2]
187 arm2Angle_t = self.mbs.GetNodeOutput(self.nArm2, variableType=exu.OutputVariableType.Coordinates_t)[2]
188 arm3Angle_t = self.mbs.GetNodeOutput(self.nArm3, variableType=exu.OutputVariableType.Coordinates_t)[2]
189
190 #finally write updated state:
191 self.state = (cartPosX, cartPosX_t, arm1Angle, arm1Angle_t, arm2Angle, arm2Angle_t, arm3Angle, arm3Angle_t)
192 #++++++++++++++++++++++++++++++++++++++++++++++++++
193
194 done = bool(
195 cartPosX < -self.x_threshold
196 or cartPosX > self.x_threshold
197 or arm1Angle < -self.theta_threshold_radians
198 or arm1Angle > self.theta_threshold_radians
199 or arm2Angle < -self.theta_threshold_radians
200 or arm2Angle > self.theta_threshold_radians
201 or arm3Angle < -self.theta_threshold_radians
202 or arm3Angle > self.theta_threshold_radians
203 )
204 return done
205
206
207 #**classFunction: OVERRIDE this function to maps the current state to mbs initial values
208 #**output: return `[initialValues, initialValues_t]`, where `initialValues` and `initialValues_t` are the ODE2 coordinate and velocity vectors of the mbs
209 def State2InitialValues(self):
210 #+++++++++++++++++++++++++++++++++++++++++++++
211 #set specific initial state:
212 (xCart, xCart_t, phiArm1, phiArm1_t, phiArm2, phiArm2_t, phiArm3, phiArm3_t) = self.state
213
214 initialValues = np.zeros(12) #model has 4*3 redundant states
215 initialValues_t = np.zeros(12)
216
217 #build redundant cordinates from self.state
218 initialValues[0] = xCart
219 initialValues[3+0] = xCart - 0.5*self.length * sin(phiArm1)
220 initialValues[3+1] = 0.5*self.length * (cos(phiArm1)-1)
221 initialValues[3+2] = phiArm1
222
223 initialValues[6+0] = xCart - self.length * sin(phiArm1) - 0.5*self.length * sin(phiArm2)
224 initialValues[6+1] = self.length * cos(phiArm1) + 0.5*self.length * cos(phiArm2) - 1.5*self.length
225 initialValues[6+2] = phiArm2
226
227 initialValues[9+0] = xCart - self.length * sin(phiArm1) - self.length * sin(phiArm2) - 0.5*self.length * sin(phiArm3)
228 initialValues[9+1] = self.length * cos(phiArm1) + self.length * cos(phiArm2) + 0.5*self.length * cos(phiArm3) - 2.5*self.length
229 initialValues[9+2] = phiArm3
230
231 initialValues_t[0] = xCart_t
232 initialValues_t[3+0] = xCart_t - phiArm1_t*0.5*self.length * cos(phiArm1)
233 initialValues_t[3+1] = -0.5*self.length * sin(phiArm1) * phiArm1_t
234 initialValues_t[3+2] = phiArm1_t
235
236 initialValues_t[6+0] = xCart_t - phiArm1_t*self.length * cos(phiArm1) - phiArm2_t*0.5*self.length * cos(phiArm2)
237 initialValues_t[6+1] = -self.length * sin(phiArm1) * phiArm1_t - 0.5*self.length * sin(phiArm2) * phiArm2_t
238 initialValues_t[6+2] = phiArm2_t
239
240 initialValues_t[9+0] = xCart_t - phiArm1_t*self.length * cos(phiArm1) - phiArm2_t*self.length * cos(phiArm2) - phiArm3_t*0.5*self.length * cos(phiArm3)
241 initialValues_t[9+1] = -self.length * sin(phiArm1) * phiArm1_t - self.length * sin(phiArm2) * phiArm2_t - 0.5*self.length * sin(phiArm3) * phiArm3_t
242 initialValues_t[9+2] = phiArm3_t
243
244 return [initialValues,initialValues_t]
245
246
247
248
249
250
251
252
253
254#%%+++++++++++++++++++++++++++++++++++++++++++++
255if __name__ == '__main__': #this is only executed when file is direct called in Python
256 import time
257
258
259 #%%++++++++++++++++++++++++++++++++++++++++++++++++++
260 #use some learning algorithm:
261 #pip install stable_baselines3
262 from stable_baselines3 import A2C
263
264
265 #create model and do reinforcement learning
266 if False: #'scalar' environment:
267 env = InvertedTriplePendulumEnv() #(thresholdFactor=2)
268 #check if model runs:
269 # env.TestModel(numberOfSteps=1000, seed=42)
270
271 #main learning task; 1e7 steps take 2-3 hours
272 model = A2C('MlpPolicy',
273 env,
274 device='cpu', #usually cpu is faster for this size of networks
275 #device='cuda', #usually cpu is faster for this size of networks
276 verbose=1)
277 ts = -time.time()
278 model.learn(total_timesteps=2000)
279 #model.learn(total_timesteps=2e7) #not sufficient ...
280 print('*** learning time total =',ts+time.time(),'***')
281
282 #save learned model
283 model.save("openAIgymTriplePendulum1e7d")
284 else:
285 #create vectorized environment, which is much faster for time
286 # consuming environments (otherwise learning algo may be the bottleneck)
287 # https://www.programcreek.com/python/example/121472/stable_baselines.common.vec_env.SubprocVecEnv
288 import torch #stable-baselines3 is based on pytorch
289 n_cores= max(1,int(os.cpu_count()/2)) #n_cores should be number of real cores (not threads)
290 #n_cores=14 #should be number of real cores (not threads)
291 torch.set_num_threads(n_cores) #seems to be ideal to match the size of subprocVecEnv
292
293 #test problem with nSteps=400 in time integration
294 #1 core: learning time total = 28.73 seconds
295 #4 core: learning time total = 8.10
296 #8 core: learning time total = 4.48
297 #14 core:learning time total = 3.77
298 #standard DummyVecEnv version: 15.14 seconds
299 print('using',n_cores,'cores')
300
301 from stable_baselines3.common.vec_env import DummyVecEnv, SubprocVecEnv
302 vecEnv = SubprocVecEnv([InvertedTriplePendulumEnv for i in range(n_cores)])
303
304
305 #main learning task; 1e7 steps take 2-3 hours
306 model = A2C('MlpPolicy',
307 vecEnv,
308 device='cpu', #usually cpu is faster for this size of networks
309 #device='cuda', #optimal with 64 SubprocVecEnv, torch.set_num_threads(1)
310 verbose=1)
311 ts = -time.time()
312 print('start learning...')
313 #model.learn(total_timesteps=50000)
314 model.learn(total_timesteps=7e7) #not sufficient ...
315 print('*** learning time total =',ts+time.time(),'***')
316
317 #save learned model
318 model.save("openAIgymTriplePendulum1e7d")
319
320 if False:
321 #%%++++++++++++++++++++++++++++++++++++++++++++++++++
322 #only load and test
323 model = A2C.load("openAIgymTriplePendulum1e7")
324 env = InvertedTriplePendulumEnv(thresholdFactor=15) #larger threshold for testing
325 solutionFile='solution/learningCoordinates.txt'
326 env.TestModel(numberOfSteps=2500, model=model, solutionFileName=solutionFile,
327 stopIfDone=False, useRenderer=False, sleepTime=0) #just compute solution file
328
329 #++++++++++++++++++++++++++++++++++++++++++++++
330 #visualize (and make animations) in exudyn:
331 from exudyn.interactive import SolutionViewer
332 env.SC.visualizationSettings.general.autoFitScene = False
333 solution = LoadSolutionFile(OutputFilePath(solutionFile))
334 SolutionViewer(env.mbs, solution) #loads solution file via name stored in mbs