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