mirror of
https://github.com/Farama-Foundation/Gymnasium.git
synced 2025-08-01 22:11:25 +00:00
119 lines
4.4 KiB
Python
119 lines
4.4 KiB
Python
"""
|
|
Classic cart-pole system implemented by Rich Sutton et al.
|
|
Copied from https://webdocs.cs.ualberta.ca/~sutton/book/code/pole.c
|
|
"""
|
|
|
|
import math
|
|
import gym
|
|
from gym import spaces
|
|
import numpy as np
|
|
|
|
class CartPoleEnv(gym.Env):
|
|
metadata = {
|
|
'render.modes': ['human', 'rgb_array'],
|
|
'video.frames_per_second' : 50
|
|
}
|
|
|
|
def __init__(self):
|
|
self.gravity = 9.8
|
|
self.masscart = 1.0
|
|
self.masspole = 0.1
|
|
self.total_mass = (self.masspole + self.masscart)
|
|
self.length = 0.5 # actually half the pole's length
|
|
self.polemass_length = (self.masspole * self.length)
|
|
self.force_mag = 10.0
|
|
self.tau = 0.02 # seconds between state updates
|
|
|
|
# Angle at which to fail the episode
|
|
self.theta_threshold_radians = 12 * 2 * math.pi / 360
|
|
self.x_threshold = 2.4
|
|
self.reset()
|
|
self.viewer = None
|
|
|
|
high = np.array([self.x_threshold, np.inf, self.theta_threshold_radians, np.inf])
|
|
self.action_space = spaces.Discrete(2)
|
|
self.observation_space = spaces.Box(-high, high)
|
|
|
|
def _step(self, action):
|
|
action = action
|
|
assert action==0 or action==1, "%r (%s) invalid"%(action, type(action))
|
|
state = self.state
|
|
x, x_dot, theta, theta_dot = state
|
|
force = self.force_mag if action==1 else -self.force_mag
|
|
costheta = math.cos(theta)
|
|
sintheta = math.sin(theta)
|
|
temp = (force + self.polemass_length * theta_dot * theta_dot * sintheta) / self.total_mass
|
|
thetaacc = (self.gravity * sintheta - costheta* temp) / (self.length * (4.0/3.0 - self.masspole * costheta * costheta / self.total_mass))
|
|
xacc = temp - self.polemass_length * thetaacc * costheta / self.total_mass
|
|
x = x + self.tau * x_dot
|
|
x_dot = x_dot + self.tau * xacc
|
|
theta = theta + self.tau * theta_dot
|
|
theta_dot = theta_dot + self.tau * thetaacc
|
|
self.state = (x,x_dot,theta,theta_dot)
|
|
done = x < -self.x_threshold \
|
|
or x > self.x_threshold \
|
|
or theta < -self.theta_threshold_radians \
|
|
or theta > self.theta_threshold_radians
|
|
done = bool(done)
|
|
reward = 1.0
|
|
return np.array(self.state), reward, done, {}
|
|
|
|
def _reset(self):
|
|
self.state = np.random.uniform(low=-0.05, high=0.05, size=(4,))
|
|
return np.array(self.state)
|
|
|
|
def _render(self, mode='human', close=False):
|
|
if close:
|
|
if self.viewer is not None:
|
|
self.viewer.close()
|
|
return
|
|
|
|
screen_width = 600
|
|
screen_height = 400
|
|
|
|
world_width = self.x_threshold*2
|
|
scale = screen_width/world_width
|
|
carty = 100 # TOP OF CART
|
|
polewidth = 10.0
|
|
polelen = scale * 1.0
|
|
cartwidth = 50.0
|
|
cartheight = 30.0
|
|
|
|
if self.viewer is None:
|
|
from gym.envs.classic_control import rendering
|
|
self.viewer = rendering.Viewer(screen_width, screen_height)
|
|
l,r,t,b = -cartwidth/2, cartwidth/2, cartheight/2, -cartheight/2
|
|
axleoffset =cartheight/4.0
|
|
cart = rendering.FilledPolygon([(l,b), (l,t), (r,t), (r,b)])
|
|
self.carttrans = rendering.Transform()
|
|
cart.add_attr(self.carttrans)
|
|
self.viewer.add_geom(cart)
|
|
l,r,t,b = -polewidth/2,polewidth/2,polelen-polewidth/2,-polewidth/2
|
|
pole = rendering.FilledPolygon([(l,b), (l,t), (r,t), (r,b)])
|
|
pole.set_color(.8,.6,.4)
|
|
self.poletrans = rendering.Transform(translation=(0, axleoffset))
|
|
pole.add_attr(self.poletrans)
|
|
pole.add_attr(self.carttrans)
|
|
self.viewer.add_geom(pole)
|
|
self.axle = rendering.make_circle(polewidth/2)
|
|
self.axle.add_attr(self.poletrans)
|
|
self.axle.add_attr(self.carttrans)
|
|
self.axle.set_color(.5,.5,.8)
|
|
self.viewer.add_geom(self.axle)
|
|
self.track = rendering.Line((0,carty), (screen_width,carty))
|
|
self.track.set_color(0,0,0)
|
|
self.viewer.add_geom(self.track)
|
|
|
|
x = self.state
|
|
cartx = x[0]*scale+screen_width/2.0 # MIDDLE OF CART
|
|
self.carttrans.set_translation(cartx, carty)
|
|
self.poletrans.set_rotation(-x[2])
|
|
|
|
self.viewer.render()
|
|
if mode == 'rgb_array':
|
|
return self.viewer.get_array()
|
|
elif mode is 'human':
|
|
pass
|
|
else:
|
|
return super(CartPoleEnv, self).render(mode=mode)
|