testGymDoublePendulumEnv.py
You can view and download this file on Github: testGymDoublePendulumEnv.py
1#+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
2# This is an EXUDYN example
3#
4# Details: This file serves as an input to testGymDoublePendulum.py
5#
6# Author: Johannes Gerstmayr
7# Date: 2022-05-18
8# Update: 2023-05-20: derive from gym.Env to ensure compatibility with newer stable-baselines3
9#
10# 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.
11#
12#+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
13
14import exudyn as exu
15from exudyn.utilities import * #includes itemInterface and rigidBodyUtilities
16import exudyn
17from math import sin, cos
18import exudyn.graphics as graphics
19import math
20
21import math
22from typing import Optional, Union
23
24import numpy as np
25
26import gym
27from gym import logger, spaces, Env
28from gym.error import DependencyNotInstalled
29
30import stable_baselines3
31useOldGym = tuple(map(int, stable_baselines3.__version__.split('.'))) <= tuple(map(int, '1.8.0'.split('.')))
32
33class DoublePendulumEnv(Env):
34
35 #metadata = {"render_modes": ["human", "rgb_array"], "render_fps": 50}
36 metadata = {"render_modes": ["human"], "render_fps": 50}
37
38 def __init__(self, thresholdFactor = 1.):
39
40 self.SC = exu.SystemContainer()
41 self.mbs = self.SC.AddSystem()
42
43
44 #%%+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
45 #+++++++++++++++++++++++++++++++++++++++++++++++++++++
46 self.gravity = 9.81
47 self.length = 1.
48 self.masscart = 1.
49 self.massarm = 0.1
50 self.total_mass = self.massarm + self.masscart
51 self.armInertia = self.length**2*0.5*self.massarm
52 self.force_mag = 10.0
53 self.stepUpdateTime = 0.02 # seconds between state updates
54
55 # Angle at which to fail the episode
56 self.theta_threshold_radians = thresholdFactor* 12 * 2 * math.pi / 360
57 self.x_threshold = thresholdFactor*2.4
58
59 high = np.array(
60 [
61 self.x_threshold * 2,
62 np.finfo(np.float32).max,
63 self.theta_threshold_radians * 2,
64 np.finfo(np.float32).max,
65 self.theta_threshold_radians * 2,
66 np.finfo(np.float32).max,
67 ],
68 dtype=np.float32,
69 )
70
71 #+++++++++++++++++++++++++++++++++++++++++++++++++++++
72 #see https://github.com/openai/gym/blob/64b4b31d8245f6972b3d37270faf69b74908a67d/gym/core.py#L16
73 #for Env:
74 self.action_space = spaces.Discrete(2)
75 self.observation_space = spaces.Box(-high, high, dtype=np.float32)
76 #+++++++++++++++++++++++++++++++++++++++++++++++++++++
77 self.state = None
78 self.rendererRunning=None
79 self.useRenderer = False #turn this on if needed
80
81 background = graphics.CheckerBoard(point= [0,0,0], normal= [0,0,1], size=4)
82
83 oGround=self.mbs.AddObject(ObjectGround(referencePosition= [0,0,0], #x-pos,y-pos,angle
84 visualization=VObjectGround(graphicsData= [background])))
85 nGround=self.mbs.AddNode(NodePointGround())
86
87 gCart = graphics.Brick(size=[0.5*self.length, 0.1*self.length, 0.1*self.length],
88 color=graphics.color.dodgerblue)
89 self.nCart = self.mbs.AddNode(Rigid2D(referenceCoordinates=[0,0,0]));
90 oCart = self.mbs.AddObject(RigidBody2D(mass=self.masscart,
91 inertia=0.1*self.masscart, #not needed
92 nodeNumber=self.nCart,
93 visualization=VObjectRigidBody2D(graphicsData= [gCart])))
94 mCartCOM = self.mbs.AddMarker(MarkerNodePosition(nodeNumber=self.nCart))
95
96 gArm1 = graphics.Brick(size=[0.1*self.length, self.length, 0.1*self.length], color=graphics.color.red)
97 self.nArm1 = self.mbs.AddNode(Rigid2D(referenceCoordinates=[0,0.5*self.length,0]));
98 oArm1 = self.mbs.AddObject(RigidBody2D(mass=self.massarm,
99 inertia=self.armInertia, #not included in original paper
100 nodeNumber=self.nArm1,
101 visualization=VObjectRigidBody2D(graphicsData= [gArm1])))
102
103 mArm1COM = self.mbs.AddMarker(MarkerNodePosition(nodeNumber=self.nArm1))
104 mArm1JointA = self.mbs.AddMarker(MarkerBodyPosition(bodyNumber=oArm1, localPosition=[0,-0.5*self.length,0]))
105 mArm1JointB = self.mbs.AddMarker(MarkerBodyPosition(bodyNumber=oArm1, localPosition=[0, 0.5*self.length,0]))
106
107 gArm2 = graphics.Brick(size=[0.1*self.length, self.length, 0.1*self.length], color=graphics.color.red)
108 self.nArm2 = self.mbs.AddNode(Rigid2D(referenceCoordinates=[0,1.5*self.length,0]));
109 oArm2 = self.mbs.AddObject(RigidBody2D(mass=self.massarm,
110 inertia=self.armInertia, #not included in original paper
111 nodeNumber=self.nArm2,
112 visualization=VObjectRigidBody2D(graphicsData= [gArm2])))
113
114 mArm2COM = self.mbs.AddMarker(MarkerNodePosition(nodeNumber=self.nArm2))
115 mArm2Joint = self.mbs.AddMarker(MarkerBodyPosition(bodyNumber=oArm2, localPosition=[0,-0.5*self.length,0]))
116
117 mCartCoordX = self.mbs.AddMarker(MarkerNodeCoordinate(nodeNumber=self.nCart, coordinate=0))
118 mCartCoordY = self.mbs.AddMarker(MarkerNodeCoordinate(nodeNumber=self.nCart, coordinate=1))
119 mGroundNode = self.mbs.AddMarker(MarkerNodeCoordinate(nodeNumber=nGround, coordinate=0))
120
121 #gravity
122 self.mbs.AddLoad(Force(markerNumber=mCartCOM, loadVector=[0,-self.masscart*self.gravity,0]))
123 self.mbs.AddLoad(Force(markerNumber=mArm1COM, loadVector=[0,-self.massarm*self.gravity,0]))
124 self.mbs.AddLoad(Force(markerNumber=mArm2COM, loadVector=[0,-self.massarm*self.gravity,0]))
125
126 #control force
127 self.lControl = self.mbs.AddLoad(LoadCoordinate(markerNumber=mCartCoordX, load=1.))
128
129 #joints and constraints:
130 self.mbs.AddObject(RevoluteJoint2D(markerNumbers=[mCartCOM, mArm1JointA]))
131 self.mbs.AddObject(RevoluteJoint2D(markerNumbers=[mArm1JointB, mArm2Joint]))
132 self.mbs.AddObject(CoordinateConstraint(markerNumbers=[mCartCoordY, mGroundNode]))
133
134
135
136
137 #%%++++++++++++++++++++++++
138 self.mbs.Assemble() #computes initial vector
139
140 self.simulationSettings = exu.SimulationSettings() #takes currently set values or default values
141
142
143 self.simulationSettings.timeIntegration.numberOfSteps = 1
144 self.simulationSettings.timeIntegration.endTime = 0 #will be overwritten in step
145 self.simulationSettings.timeIntegration.verboseMode = 0
146 self.simulationSettings.solution.file.write = False
147 #self.simulationSettings.timeIntegration.realtime.active = True
148
149 self.simulationSettings.timeIntegration.newton.useModifiedNewton = True
150
151 self.SC.visualizationSettings.view0.scene.drawWorldBasis=True
152 self.SC.visualizationSettings.general.graphicsUpdateInterval = 0.01 #50Hz
153
154 self.simulationSettings.solution.file.information = "Open AI gym"
155
156 self.dynamicSolver = exudyn.MainSolverImplicitSecondOrder()
157 self.dynamicSolver.InitializeSolver(self.mbs, self.simulationSettings)
158 self.dynamicSolver.SolveSteps(self.mbs, self.simulationSettings) #to initialize all data
159
160 def integrateStep(self, force):
161 #exudyn simulation part
162 #index 2 solver
163 self.mbs.SetLoadParameter(self.lControl, 'load', force)
164
165 #progress integration time
166 currentTime = self.simulationSettings.timeIntegration.endTime
167 self.simulationSettings.timeIntegration.startTime = currentTime
168 self.simulationSettings.timeIntegration.endTime = currentTime+self.stepUpdateTime
169
170 self.dynamicSolver.InitializeSolverInitialConditions(self.mbs, self.simulationSettings)
171 self.dynamicSolver.SolveSteps(self.mbs, self.simulationSettings)
172 currentState = self.mbs.systemData.GetSystemState() #get current values
173 self.mbs.systemData.SetSystemState(systemStateList=currentState,
174 configuration = exu.ConfigurationType.Initial)
175 self.mbs.systemData.SetODE2Coordinates_tt(coordinates = self.mbs.systemData.GetODE2Coordinates_tt(),
176 configuration = exu.ConfigurationType.Initial)
177
178
179
180 def step(self, action):
181 err_msg = f"{action!r} ({type(action)}) invalid"
182 assert self.action_space.contains(action), err_msg
183 assert self.state is not None, "Call reset before using step method."
184 self.setState2InitialValues()
185
186 force = self.force_mag if action == 1 else -self.force_mag
187
188 #++++++++++++++++++++++++++++++++++++++++++++++++++
189 #++++++++++++++++++++++++++++++++++++++++++++++++++
190 self.integrateStep(force)
191 #+++++++++++++++++++++++++
192 #compute some output:
193 cartPosX = self.mbs.GetNodeOutput(self.nCart, variableType=exu.OutputVariableType.Coordinates)[0]
194 arm1Angle = self.mbs.GetNodeOutput(self.nArm1, variableType=exu.OutputVariableType.Coordinates)[2]
195 arm2Angle = self.mbs.GetNodeOutput(self.nArm2, variableType=exu.OutputVariableType.Coordinates)[2]
196 cartPosX_t = self.mbs.GetNodeOutput(self.nCart, variableType=exu.OutputVariableType.Coordinates_t)[0]
197 arm1Angle_t = self.mbs.GetNodeOutput(self.nArm1, variableType=exu.OutputVariableType.Coordinates_t)[2]
198 arm2Angle_t = self.mbs.GetNodeOutput(self.nArm2, variableType=exu.OutputVariableType.Coordinates_t)[2]
199
200 #finally write updated state:
201 self.state = (cartPosX, cartPosX_t, arm1Angle, arm1Angle_t, arm2Angle, arm2Angle_t)
202 #++++++++++++++++++++++++++++++++++++++++++++++++++
203 #++++++++++++++++++++++++++++++++++++++++++++++++++
204
205 done = bool(
206 cartPosX < -self.x_threshold
207 or cartPosX > self.x_threshold
208 or arm1Angle < -self.theta_threshold_radians
209 or arm1Angle > self.theta_threshold_radians
210 or arm2Angle < -self.theta_threshold_radians
211 or arm2Angle > self.theta_threshold_radians
212 )
213
214 if not done:
215 reward = 1.0
216 elif self.steps_beyond_done is None:
217 # Arm1 just fell!
218 self.steps_beyond_done = 0
219 reward = 1.0
220 else:
221 if self.steps_beyond_done == 0:
222 logger.warn(
223 "You are calling 'step()' even though this "
224 "environment has already returned done = True. You "
225 "should always call 'reset()' once you receive 'done = "
226 "True' -- any further steps are undefined behavior."
227 )
228 self.steps_beyond_done += 1
229 reward = 0.0
230
231 info = {}
232 terminated, truncated = done, False # since stable-baselines3 > 1.8.0 implementations terminated and truncated
233 if useOldGym:
234 return np.array(self.state, dtype=np.float32), reward, terminated, info
235 else:
236 return np.array(self.state, dtype=np.float32), reward, terminated, truncated, info
237
238 def setState2InitialValues(self):
239 #+++++++++++++++++++++++++++++++++++++++++++++
240 #set specific initial state:
241 (xCart, xCart_t, phiArm1, phiArm1_t, phiArm2, phiArm2_t) = self.state
242
243 initialValues = np.zeros(9) #model has 3*3 = 9 redundant states
244 initialValues_t = np.zeros(9)
245
246 #build redundant cordinates from self.state
247 initialValues[0] = xCart
248 initialValues[3+0] = xCart - 0.5*self.length * sin(phiArm1)
249 initialValues[3+1] = 0.5*self.length * (cos(phiArm1)-1)
250 initialValues[3+2] = phiArm1
251 initialValues[6+0] = xCart - 1.*self.length * sin(phiArm1) - 0.5*self.length * sin(phiArm2)
252 initialValues[6+1] = self.length * cos(phiArm1) + 0.5*self.length * cos(phiArm2) - 1.5*self.length
253 initialValues[6+2] = phiArm2
254
255 initialValues_t[0] = xCart_t
256 initialValues_t[3+0] = xCart_t - phiArm1_t*0.5*self.length * cos(phiArm1)
257 initialValues_t[3+1] = -0.5*self.length * sin(phiArm1) * phiArm1_t
258 initialValues_t[3+2] = phiArm1_t
259 initialValues_t[6+0] = xCart_t - phiArm1_t*1.*self.length * cos(phiArm1) - phiArm2_t*0.5*self.length * cos(phiArm2)
260 initialValues_t[6+1] = -self.length * sin(phiArm1) * phiArm1_t - 0.5*self.length * sin(phiArm2) * phiArm2_t
261 initialValues_t[6+2] = phiArm2_t
262
263 self.mbs.systemData.SetODE2Coordinates(initialValues, exu.ConfigurationType.Initial)
264 self.mbs.systemData.SetODE2Coordinates_t(initialValues_t, exu.ConfigurationType.Initial)
265
266
267 def reset(
268 self,
269 *,
270 seed: Optional[int] = None,
271 return_info: bool = False,
272 options: Optional[dict] = None,
273 ):
274 #super().reset(seed=seed)
275 self.state = np.random.uniform(low=-0.05, high=0.05, size=(6,))
276 self.steps_beyond_done = None
277
278
279 #+++++++++++++++++++++++++++++++++++++++++++++
280 #set state into initial values:
281 self.setState2InitialValues()
282
283 self.simulationSettings.timeIntegration.endTime = 0
284 self.dynamicSolver.InitializeSolver(self.mbs, self.simulationSettings) #needed to update initial conditions
285
286 #+++++++++++++++++++++++++++++++++++++++++++++
287
288 if not return_info and useOldGym:
289 return np.array(self.state, dtype=np.float32)
290 else:
291 return np.array(self.state, dtype=np.float32), {}
292
293 def render(self, mode="human"):
294 if self.rendererRunning==None and self.useRenderer:
295 SC.renderer.Start()
296 self.rendererRunning = True
297
298 def close(self):
299 self.dynamicSolver.FinalizeSolver(self.mbs, self.simulationSettings)
300 if self.rendererRunning==True:
301 # SC.renderer.DoIdleTasks()
302 SC.renderer.Stop() #safely close rendering window!
303
304
305
306# #+++++++++++++++++++++++++++++++++++++++++++++
307# #reset:
308# self.mbs.systemData.SetODE2Coordinates(initialValues, exu.ConfigurationType.Initial)
309# self.mbs.systemData.SetODE2Coordinates_t(initialValues, exu.ConfigurationType.Initial)
310# self.mbs.systemData.SetODE2Coordinates_tt(initialValues, exu.ConfigurationType.Initial)