mirror of
https://github.com/Farama-Foundation/Gymnasium.git
synced 2025-08-01 06:07:08 +00:00
Refactor classic control rendering to use pygame (#2599)
* refactor classic control rendering to use pygame * Rebase from upstream * Revert deletion from rebasing * Remove pyglet
This commit is contained in:
Andrew Tan Jin Shen
committed by
GitHub
GitHub
parent
d21dec1829
commit
6eba48db95
@@ -2,6 +2,8 @@
|
||||
from typing import Optional
|
||||
|
||||
import numpy as np
|
||||
import pygame
|
||||
from pygame import gfxdraw
|
||||
from numpy import sin, cos, pi
|
||||
|
||||
from gym import core, spaces
|
||||
@@ -146,6 +148,8 @@ class AcrobotEnv(core.Env):
|
||||
|
||||
torque_noise_max = 0.0
|
||||
|
||||
SCREEN_DIM = 500
|
||||
|
||||
#: use dynamics equations from the nips paper or the book
|
||||
book_or_nips = "book"
|
||||
action_arrow = None
|
||||
@@ -153,7 +157,8 @@ class AcrobotEnv(core.Env):
|
||||
actions_num = 3
|
||||
|
||||
def __init__(self):
|
||||
self.viewer = None
|
||||
self.screen = None
|
||||
self.isopen = True
|
||||
high = np.array(
|
||||
[1.0, 1.0, 1.0, 1.0, self.MAX_VEL_1, self.MAX_VEL_2], dtype=np.float32
|
||||
)
|
||||
@@ -259,46 +264,75 @@ class AcrobotEnv(core.Env):
|
||||
return (dtheta1, dtheta2, ddtheta1, ddtheta2, 0.0)
|
||||
|
||||
def render(self, mode="human"):
|
||||
from gym.utils import pyglet_rendering
|
||||
|
||||
if self.screen is None:
|
||||
pygame.init()
|
||||
self.screen = pygame.display.set_mode((self.SCREEN_DIM, self.SCREEN_DIM))
|
||||
self.surf = pygame.Surface((self.SCREEN_DIM, self.SCREEN_DIM))
|
||||
self.surf.fill((255, 255, 255))
|
||||
s = self.state
|
||||
|
||||
if self.viewer is None:
|
||||
self.viewer = pyglet_rendering.Viewer(500, 500)
|
||||
bound = self.LINK_LENGTH_1 + self.LINK_LENGTH_2 + 0.2 # 2.2 for default
|
||||
self.viewer.set_bounds(-bound, bound, -bound, bound)
|
||||
bound = self.LINK_LENGTH_1 + self.LINK_LENGTH_2 + 0.2 # 2.2 for default
|
||||
scale = self.SCREEN_DIM / (bound * 2)
|
||||
offset = self.SCREEN_DIM / 2
|
||||
|
||||
if s is None:
|
||||
return None
|
||||
|
||||
p1 = [-self.LINK_LENGTH_1 * cos(s[0]), self.LINK_LENGTH_1 * sin(s[0])]
|
||||
p1 = [
|
||||
-self.LINK_LENGTH_1 * cos(s[0]) * scale,
|
||||
self.LINK_LENGTH_1 * sin(s[0]) * scale,
|
||||
]
|
||||
|
||||
p2 = [
|
||||
p1[0] - self.LINK_LENGTH_2 * cos(s[0] + s[1]),
|
||||
p1[1] + self.LINK_LENGTH_2 * sin(s[0] + s[1]),
|
||||
p1[0] - self.LINK_LENGTH_2 * cos(s[0] + s[1]) * scale,
|
||||
p1[1] + self.LINK_LENGTH_2 * sin(s[0] + s[1]) * scale,
|
||||
]
|
||||
|
||||
xys = np.array([[0, 0], p1, p2])[:, ::-1]
|
||||
thetas = [s[0] - pi / 2, s[0] + s[1] - pi / 2]
|
||||
link_lengths = [self.LINK_LENGTH_1, self.LINK_LENGTH_2]
|
||||
link_lengths = [self.LINK_LENGTH_1 * scale, self.LINK_LENGTH_2 * scale]
|
||||
|
||||
pygame.draw.line(
|
||||
self.surf,
|
||||
start_pos=(-2.2 * scale + offset, 1 * scale + offset),
|
||||
end_pos=(2.2 * scale + offset, 1 * scale + offset),
|
||||
color=(0, 0, 0),
|
||||
)
|
||||
|
||||
self.viewer.draw_line((-2.2, 1), (2.2, 1))
|
||||
for ((x, y), th, llen) in zip(xys, thetas, link_lengths):
|
||||
l, r, t, b = 0, llen, 0.1, -0.1
|
||||
jtransform = pyglet_rendering.Transform(rotation=th, translation=(x, y))
|
||||
link = self.viewer.draw_polygon([(l, b), (l, t), (r, t), (r, b)])
|
||||
link.add_attr(jtransform)
|
||||
link.set_color(0, 0.8, 0.8)
|
||||
circ = self.viewer.draw_circle(0.1)
|
||||
circ.set_color(0.8, 0.8, 0)
|
||||
circ.add_attr(jtransform)
|
||||
x = x + offset
|
||||
y = y + offset
|
||||
l, r, t, b = 0, llen, 0.1 * scale, -0.1 * scale
|
||||
coords = [(l, b), (l, t), (r, t), (r, b)]
|
||||
transformed_coords = []
|
||||
for coord in coords:
|
||||
coord = pygame.math.Vector2(coord).rotate_rad(th)
|
||||
coord = (coord[0] + x, coord[1] + y)
|
||||
transformed_coords.append(coord)
|
||||
gfxdraw.aapolygon(self.surf, transformed_coords, (0, 204, 204))
|
||||
gfxdraw.filled_polygon(self.surf, transformed_coords, (0, 204, 204))
|
||||
|
||||
return self.viewer.render(return_rgb_array=mode == "rgb_array")
|
||||
gfxdraw.aacircle(self.surf, int(x), int(y), int(0.1 * scale), (204, 204, 0))
|
||||
gfxdraw.filled_circle(
|
||||
self.surf, int(x), int(y), int(0.1 * scale), (204, 204, 0)
|
||||
)
|
||||
|
||||
self.surf = pygame.transform.flip(self.surf, False, True)
|
||||
self.screen.blit(self.surf, (0, 0))
|
||||
if mode == "human":
|
||||
pygame.display.flip()
|
||||
|
||||
if mode == "rgb_array":
|
||||
return np.transpose(
|
||||
np.array(pygame.surfarray.pixels3d(self.screen)), axes=(1, 0, 2)
|
||||
)
|
||||
else:
|
||||
return self.isopen
|
||||
|
||||
def close(self):
|
||||
if self.viewer:
|
||||
self.viewer.close()
|
||||
self.viewer = None
|
||||
if self.screen is not None:
|
||||
pygame.quit()
|
||||
self.isopen = False
|
||||
|
||||
|
||||
def wrap(x, m, M):
|
||||
|
||||
@@ -6,10 +6,13 @@ permalink: https://perma.cc/C9ZM-652R
|
||||
import math
|
||||
from typing import Optional, Union
|
||||
|
||||
import numpy as np
|
||||
import pygame
|
||||
from pygame import gfxdraw
|
||||
|
||||
import gym
|
||||
from gym import spaces, logger
|
||||
from gym.utils import seeding
|
||||
import numpy as np
|
||||
|
||||
|
||||
class CartPoleEnv(gym.Env[np.ndarray, Union[int, np.ndarray]]):
|
||||
@@ -101,7 +104,8 @@ class CartPoleEnv(gym.Env[np.ndarray, Union[int, np.ndarray]]):
|
||||
self.action_space = spaces.Discrete(2)
|
||||
self.observation_space = spaces.Box(-high, high, dtype=np.float32)
|
||||
|
||||
self.viewer = None
|
||||
self.screen = None
|
||||
self.isopen = True
|
||||
self.state = None
|
||||
|
||||
self.steps_beyond_done = None
|
||||
@@ -185,66 +189,73 @@ class CartPoleEnv(gym.Env[np.ndarray, Union[int, np.ndarray]]):
|
||||
|
||||
world_width = self.x_threshold * 2
|
||||
scale = screen_width / world_width
|
||||
carty = 100 # TOP OF CART
|
||||
polewidth = 10.0
|
||||
polelen = scale * (2 * self.length)
|
||||
cartwidth = 50.0
|
||||
cartheight = 30.0
|
||||
|
||||
if self.viewer is None:
|
||||
from gym.utils import pyglet_rendering
|
||||
x = self.state
|
||||
|
||||
self.viewer = pyglet_rendering.Viewer(screen_width, screen_height)
|
||||
l, r, t, b = -cartwidth / 2, cartwidth / 2, cartheight / 2, -cartheight / 2
|
||||
axleoffset = cartheight / 4.0
|
||||
cart = pyglet_rendering.FilledPolygon([(l, b), (l, t), (r, t), (r, b)])
|
||||
self.carttrans = pyglet_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 = pyglet_rendering.FilledPolygon([(l, b), (l, t), (r, t), (r, b)])
|
||||
pole.set_color(0.8, 0.6, 0.4)
|
||||
self.poletrans = pyglet_rendering.Transform(translation=(0, axleoffset))
|
||||
pole.add_attr(self.poletrans)
|
||||
pole.add_attr(self.carttrans)
|
||||
self.viewer.add_geom(pole)
|
||||
self.axle = pyglet_rendering.make_circle(polewidth / 2)
|
||||
self.axle.add_attr(self.poletrans)
|
||||
self.axle.add_attr(self.carttrans)
|
||||
self.axle.set_color(0.5, 0.5, 0.8)
|
||||
self.viewer.add_geom(self.axle)
|
||||
self.track = pyglet_rendering.Line((0, carty), (screen_width, carty))
|
||||
self.track.set_color(0, 0, 0)
|
||||
self.viewer.add_geom(self.track)
|
||||
if self.screen is None:
|
||||
pygame.init()
|
||||
self.screen = pygame.display.set_mode((screen_width, screen_height))
|
||||
self.surf = pygame.Surface((screen_width, screen_height))
|
||||
self.surf.fill((255, 255, 255))
|
||||
|
||||
self._pole_geom = pole
|
||||
l, r, t, b = -cartwidth / 2, cartwidth / 2, cartheight / 2, -cartheight / 2
|
||||
axleoffset = cartheight / 4.0
|
||||
cartx = x[0] * scale + screen_width / 2.0 # MIDDLE OF CART
|
||||
carty = 100 # TOP OF CART
|
||||
cart_coords = [(l, b), (l, t), (r, t), (r, b)]
|
||||
cart_coords = [(c[0] + cartx, c[1] + carty) for c in cart_coords]
|
||||
gfxdraw.aapolygon(self.surf, cart_coords, (0, 0, 0))
|
||||
gfxdraw.filled_polygon(self.surf, cart_coords, (0, 0, 0))
|
||||
|
||||
if self.state is None:
|
||||
return None
|
||||
|
||||
# Edit the pole polygon vertex
|
||||
pole = self._pole_geom
|
||||
l, r, t, b = (
|
||||
-polewidth / 2,
|
||||
polewidth / 2,
|
||||
polelen - polewidth / 2,
|
||||
-polewidth / 2,
|
||||
)
|
||||
pole.v = [(l, b), (l, t), (r, t), (r, b)]
|
||||
|
||||
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])
|
||||
pole_coords = []
|
||||
for coord in [(l, b), (l, t), (r, t), (r, b)]:
|
||||
coord = pygame.math.Vector2(coord).rotate_rad(-x[2])
|
||||
coord = (coord[0] + cartx, coord[1] + carty + axleoffset)
|
||||
pole_coords.append(coord)
|
||||
gfxdraw.aapolygon(self.surf, pole_coords, (202, 152, 101))
|
||||
gfxdraw.filled_polygon(self.surf, pole_coords, (202, 152, 101))
|
||||
|
||||
return self.viewer.render(return_rgb_array=mode == "rgb_array")
|
||||
gfxdraw.aacircle(
|
||||
self.surf,
|
||||
int(cartx),
|
||||
int(carty + axleoffset),
|
||||
int(polewidth / 2),
|
||||
(129, 132, 203),
|
||||
)
|
||||
gfxdraw.filled_circle(
|
||||
self.surf,
|
||||
int(cartx),
|
||||
int(carty + axleoffset),
|
||||
int(polewidth / 2),
|
||||
(129, 132, 203),
|
||||
)
|
||||
|
||||
gfxdraw.hline(self.surf, 0, screen_width, carty, (0, 0, 0))
|
||||
|
||||
self.surf = pygame.transform.flip(self.surf, False, True)
|
||||
self.screen.blit(self.surf, (0, 0))
|
||||
if mode == "human":
|
||||
pygame.display.flip()
|
||||
|
||||
if mode == "rgb_array":
|
||||
return np.transpose(
|
||||
np.array(pygame.surfarray.pixels3d(self.screen)), axes=(1, 0, 2)
|
||||
)
|
||||
else:
|
||||
return self.isopen
|
||||
|
||||
def close(self):
|
||||
if self.viewer:
|
||||
self.viewer.close()
|
||||
self.viewer = None
|
||||
if self.screen is not None:
|
||||
pygame.quit()
|
||||
self.isopen = False
|
||||
|
||||
@@ -17,6 +17,8 @@ import math
|
||||
from typing import Optional
|
||||
|
||||
import numpy as np
|
||||
import pygame
|
||||
from pygame import gfxdraw
|
||||
|
||||
import gym
|
||||
from gym import spaces
|
||||
@@ -91,7 +93,8 @@ class Continuous_MountainCarEnv(gym.Env):
|
||||
[self.max_position, self.max_speed], dtype=np.float32
|
||||
)
|
||||
|
||||
self.viewer = None
|
||||
self.screen = None
|
||||
self.isopen = True
|
||||
|
||||
self.action_space = spaces.Box(
|
||||
low=self.min_action, high=self.max_action, shape=(1,), dtype=np.float32
|
||||
@@ -155,61 +158,79 @@ class Continuous_MountainCarEnv(gym.Env):
|
||||
scale = screen_width / world_width
|
||||
carwidth = 40
|
||||
carheight = 20
|
||||
|
||||
if self.viewer is None:
|
||||
from gym.utils import pyglet_rendering
|
||||
|
||||
self.viewer = pyglet_rendering.Viewer(screen_width, screen_height)
|
||||
xs = np.linspace(self.min_position, self.max_position, 100)
|
||||
ys = self._height(xs)
|
||||
xys = list(zip((xs - self.min_position) * scale, ys * scale))
|
||||
|
||||
self.track = pyglet_rendering.make_polyline(xys)
|
||||
self.track.set_linewidth(4)
|
||||
self.viewer.add_geom(self.track)
|
||||
|
||||
clearance = 10
|
||||
|
||||
l, r, t, b = -carwidth / 2, carwidth / 2, carheight, 0
|
||||
car = pyglet_rendering.FilledPolygon([(l, b), (l, t), (r, t), (r, b)])
|
||||
car.add_attr(pyglet_rendering.Transform(translation=(0, clearance)))
|
||||
self.cartrans = pyglet_rendering.Transform()
|
||||
car.add_attr(self.cartrans)
|
||||
self.viewer.add_geom(car)
|
||||
frontwheel = pyglet_rendering.make_circle(carheight / 2.5)
|
||||
frontwheel.set_color(0.5, 0.5, 0.5)
|
||||
frontwheel.add_attr(
|
||||
pyglet_rendering.Transform(translation=(carwidth / 4, clearance))
|
||||
)
|
||||
frontwheel.add_attr(self.cartrans)
|
||||
self.viewer.add_geom(frontwheel)
|
||||
backwheel = pyglet_rendering.make_circle(carheight / 2.5)
|
||||
backwheel.add_attr(
|
||||
pyglet_rendering.Transform(translation=(-carwidth / 4, clearance))
|
||||
)
|
||||
backwheel.add_attr(self.cartrans)
|
||||
backwheel.set_color(0.5, 0.5, 0.5)
|
||||
self.viewer.add_geom(backwheel)
|
||||
flagx = (self.goal_position - self.min_position) * scale
|
||||
flagy1 = self._height(self.goal_position) * scale
|
||||
flagy2 = flagy1 + 50
|
||||
flagpole = pyglet_rendering.Line((flagx, flagy1), (flagx, flagy2))
|
||||
self.viewer.add_geom(flagpole)
|
||||
flag = pyglet_rendering.FilledPolygon(
|
||||
[(flagx, flagy2), (flagx, flagy2 - 10), (flagx + 25, flagy2 - 5)]
|
||||
)
|
||||
flag.set_color(0.8, 0.8, 0)
|
||||
self.viewer.add_geom(flag)
|
||||
if self.screen is None:
|
||||
pygame.init()
|
||||
self.screen = pygame.display.set_mode((screen_width, screen_height))
|
||||
self.surf = pygame.Surface((screen_width, screen_height))
|
||||
self.surf.fill((255, 255, 255))
|
||||
|
||||
pos = self.state[0]
|
||||
self.cartrans.set_translation(
|
||||
(pos - self.min_position) * scale, self._height(pos) * scale
|
||||
)
|
||||
self.cartrans.set_rotation(math.cos(3 * pos))
|
||||
|
||||
return self.viewer.render(return_rgb_array=mode == "rgb_array")
|
||||
xs = np.linspace(self.min_position, self.max_position, 100)
|
||||
ys = self._height(xs)
|
||||
xys = list(zip((xs - self.min_position) * scale, ys * scale))
|
||||
|
||||
pygame.draw.aalines(self.surf, points=xys, closed=False, color=(0, 0, 0))
|
||||
|
||||
clearance = 10
|
||||
|
||||
l, r, t, b = -carwidth / 2, carwidth / 2, carheight, 0
|
||||
coords = []
|
||||
for c in [(l, b), (l, t), (r, t), (r, b)]:
|
||||
c = pygame.math.Vector2(c).rotate_rad(math.cos(3 * pos))
|
||||
coords.append(
|
||||
(
|
||||
c[0] + (pos - self.min_position) * scale,
|
||||
c[1] + clearance + self._height(pos) * scale,
|
||||
)
|
||||
)
|
||||
|
||||
gfxdraw.aapolygon(self.surf, coords, (0, 0, 0))
|
||||
gfxdraw.filled_polygon(self.surf, coords, (0, 0, 0))
|
||||
|
||||
for c in [(carwidth / 4, 0), (-carwidth / 4, 0)]:
|
||||
c = pygame.math.Vector2(c).rotate_rad(math.cos(3 * pos))
|
||||
wheel = (
|
||||
int(c[0] + (pos - self.min_position) * scale),
|
||||
int(c[1] + clearance + self._height(pos) * scale),
|
||||
)
|
||||
|
||||
gfxdraw.aacircle(
|
||||
self.surf, wheel[0], wheel[1], int(carheight / 2.5), (128, 128, 128)
|
||||
)
|
||||
gfxdraw.filled_circle(
|
||||
self.surf, wheel[0], wheel[1], int(carheight / 2.5), (128, 128, 128)
|
||||
)
|
||||
|
||||
flagx = int((self.goal_position - self.min_position) * scale)
|
||||
flagy1 = int(self._height(self.goal_position) * scale)
|
||||
flagy2 = flagy1 + 50
|
||||
gfxdraw.vline(self.surf, flagx, flagy1, flagy2, (0, 0, 0))
|
||||
|
||||
gfxdraw.aapolygon(
|
||||
self.surf,
|
||||
[(flagx, flagy2), (flagx, flagy2 - 10), (flagx + 25, flagy2 - 5)],
|
||||
(204, 204, 0),
|
||||
)
|
||||
gfxdraw.filled_polygon(
|
||||
self.surf,
|
||||
[(flagx, flagy2), (flagx, flagy2 - 10), (flagx + 25, flagy2 - 5)],
|
||||
(204, 204, 0),
|
||||
)
|
||||
|
||||
self.surf = pygame.transform.flip(self.surf, False, True)
|
||||
self.screen.blit(self.surf, (0, 0))
|
||||
if mode == "human":
|
||||
pygame.display.flip()
|
||||
|
||||
if mode == "rgb_array":
|
||||
return np.transpose(
|
||||
np.array(pygame.surfarray.pixels3d(self.screen)), axes=(1, 0, 2)
|
||||
)
|
||||
else:
|
||||
return self.isopen
|
||||
|
||||
def close(self):
|
||||
if self.viewer:
|
||||
self.viewer.close()
|
||||
self.viewer = None
|
||||
if self.screen is not None:
|
||||
pygame.quit()
|
||||
self.isopen = False
|
||||
|
||||
@@ -6,6 +6,8 @@ import math
|
||||
from typing import Optional
|
||||
|
||||
import numpy as np
|
||||
import pygame
|
||||
from pygame import gfxdraw
|
||||
|
||||
import gym
|
||||
from gym import spaces
|
||||
@@ -69,7 +71,8 @@ class MountainCarEnv(gym.Env):
|
||||
self.low = np.array([self.min_position, -self.max_speed], dtype=np.float32)
|
||||
self.high = np.array([self.max_position, self.max_speed], dtype=np.float32)
|
||||
|
||||
self.viewer = None
|
||||
self.screen = None
|
||||
self.isopen = True
|
||||
|
||||
self.action_space = spaces.Discrete(3)
|
||||
self.observation_space = spaces.Box(self.low, self.high, dtype=np.float32)
|
||||
@@ -118,65 +121,83 @@ class MountainCarEnv(gym.Env):
|
||||
scale = screen_width / world_width
|
||||
carwidth = 40
|
||||
carheight = 20
|
||||
|
||||
if self.viewer is None:
|
||||
from gym.utils import pyglet_rendering
|
||||
|
||||
self.viewer = pyglet_rendering.Viewer(screen_width, screen_height)
|
||||
xs = np.linspace(self.min_position, self.max_position, 100)
|
||||
ys = self._height(xs)
|
||||
xys = list(zip((xs - self.min_position) * scale, ys * scale))
|
||||
|
||||
self.track = pyglet_rendering.make_polyline(xys)
|
||||
self.track.set_linewidth(4)
|
||||
self.viewer.add_geom(self.track)
|
||||
|
||||
clearance = 10
|
||||
|
||||
l, r, t, b = -carwidth / 2, carwidth / 2, carheight, 0
|
||||
car = pyglet_rendering.FilledPolygon([(l, b), (l, t), (r, t), (r, b)])
|
||||
car.add_attr(pyglet_rendering.Transform(translation=(0, clearance)))
|
||||
self.cartrans = pyglet_rendering.Transform()
|
||||
car.add_attr(self.cartrans)
|
||||
self.viewer.add_geom(car)
|
||||
frontwheel = pyglet_rendering.make_circle(carheight / 2.5)
|
||||
frontwheel.set_color(0.5, 0.5, 0.5)
|
||||
frontwheel.add_attr(
|
||||
pyglet_rendering.Transform(translation=(carwidth / 4, clearance))
|
||||
)
|
||||
frontwheel.add_attr(self.cartrans)
|
||||
self.viewer.add_geom(frontwheel)
|
||||
backwheel = pyglet_rendering.make_circle(carheight / 2.5)
|
||||
backwheel.add_attr(
|
||||
pyglet_rendering.Transform(translation=(-carwidth / 4, clearance))
|
||||
)
|
||||
backwheel.add_attr(self.cartrans)
|
||||
backwheel.set_color(0.5, 0.5, 0.5)
|
||||
self.viewer.add_geom(backwheel)
|
||||
flagx = (self.goal_position - self.min_position) * scale
|
||||
flagy1 = self._height(self.goal_position) * scale
|
||||
flagy2 = flagy1 + 50
|
||||
flagpole = pyglet_rendering.Line((flagx, flagy1), (flagx, flagy2))
|
||||
self.viewer.add_geom(flagpole)
|
||||
flag = pyglet_rendering.FilledPolygon(
|
||||
[(flagx, flagy2), (flagx, flagy2 - 10), (flagx + 25, flagy2 - 5)]
|
||||
)
|
||||
flag.set_color(0.8, 0.8, 0)
|
||||
self.viewer.add_geom(flag)
|
||||
if self.screen is None:
|
||||
pygame.init()
|
||||
self.screen = pygame.display.set_mode((screen_width, screen_height))
|
||||
self.surf = pygame.Surface((screen_width, screen_height))
|
||||
self.surf.fill((255, 255, 255))
|
||||
|
||||
pos = self.state[0]
|
||||
self.cartrans.set_translation(
|
||||
(pos - self.min_position) * scale, self._height(pos) * scale
|
||||
)
|
||||
self.cartrans.set_rotation(math.cos(3 * pos))
|
||||
|
||||
return self.viewer.render(return_rgb_array=mode == "rgb_array")
|
||||
xs = np.linspace(self.min_position, self.max_position, 100)
|
||||
ys = self._height(xs)
|
||||
xys = list(zip((xs - self.min_position) * scale, ys * scale))
|
||||
|
||||
pygame.draw.aalines(self.surf, points=xys, closed=False, color=(0, 0, 0))
|
||||
|
||||
clearance = 10
|
||||
|
||||
l, r, t, b = -carwidth / 2, carwidth / 2, carheight, 0
|
||||
coords = []
|
||||
for c in [(l, b), (l, t), (r, t), (r, b)]:
|
||||
c = pygame.math.Vector2(c).rotate_rad(math.cos(3 * pos))
|
||||
coords.append(
|
||||
(
|
||||
c[0] + (pos - self.min_position) * scale,
|
||||
c[1] + clearance + self._height(pos) * scale,
|
||||
)
|
||||
)
|
||||
|
||||
gfxdraw.aapolygon(self.surf, coords, (0, 0, 0))
|
||||
gfxdraw.filled_polygon(self.surf, coords, (0, 0, 0))
|
||||
|
||||
for c in [(carwidth / 4, 0), (-carwidth / 4, 0)]:
|
||||
c = pygame.math.Vector2(c).rotate_rad(math.cos(3 * pos))
|
||||
wheel = (
|
||||
int(c[0] + (pos - self.min_position) * scale),
|
||||
int(c[1] + clearance + self._height(pos) * scale),
|
||||
)
|
||||
|
||||
gfxdraw.aacircle(
|
||||
self.surf, wheel[0], wheel[1], int(carheight / 2.5), (128, 128, 128)
|
||||
)
|
||||
gfxdraw.filled_circle(
|
||||
self.surf, wheel[0], wheel[1], int(carheight / 2.5), (128, 128, 128)
|
||||
)
|
||||
|
||||
flagx = int((self.goal_position - self.min_position) * scale)
|
||||
flagy1 = int(self._height(self.goal_position) * scale)
|
||||
flagy2 = flagy1 + 50
|
||||
gfxdraw.vline(self.surf, flagx, flagy1, flagy2, (0, 0, 0))
|
||||
|
||||
gfxdraw.aapolygon(
|
||||
self.surf,
|
||||
[(flagx, flagy2), (flagx, flagy2 - 10), (flagx + 25, flagy2 - 5)],
|
||||
(204, 204, 0),
|
||||
)
|
||||
gfxdraw.filled_polygon(
|
||||
self.surf,
|
||||
[(flagx, flagy2), (flagx, flagy2 - 10), (flagx + 25, flagy2 - 5)],
|
||||
(204, 204, 0),
|
||||
)
|
||||
|
||||
self.surf = pygame.transform.flip(self.surf, False, True)
|
||||
self.screen.blit(self.surf, (0, 0))
|
||||
if mode == "human":
|
||||
pygame.display.flip()
|
||||
|
||||
if mode == "rgb_array":
|
||||
return np.transpose(
|
||||
np.array(pygame.surfarray.pixels3d(self.screen)), axes=(1, 0, 2)
|
||||
)
|
||||
else:
|
||||
return self.isopen
|
||||
|
||||
def get_keys_to_action(self):
|
||||
# Control with left and right arrow keys.
|
||||
return {(): 1, (276,): 0, (275,): 2, (275, 276): 1}
|
||||
|
||||
def close(self):
|
||||
if self.viewer:
|
||||
self.viewer.close()
|
||||
self.viewer = None
|
||||
if self.screen is not None:
|
||||
pygame.quit()
|
||||
self.isopen = False
|
||||
|
||||
@@ -1,12 +1,15 @@
|
||||
__credits__ = ["Carlos Luis"]
|
||||
|
||||
from typing import Optional
|
||||
from os import path
|
||||
|
||||
import numpy as np
|
||||
import pygame
|
||||
from pygame import gfxdraw
|
||||
|
||||
import gym
|
||||
from gym import spaces
|
||||
from gym.utils import seeding
|
||||
import numpy as np
|
||||
from os import path
|
||||
|
||||
|
||||
class PendulumEnv(gym.Env):
|
||||
@@ -82,7 +85,10 @@ class PendulumEnv(gym.Env):
|
||||
self.g = g
|
||||
self.m = 1.0
|
||||
self.l = 1.0
|
||||
self.viewer = None
|
||||
self.screen = None
|
||||
self.isopen = True
|
||||
|
||||
self.screen_dim = 500
|
||||
|
||||
high = np.array([1.0, 1.0, self.max_speed], dtype=np.float32)
|
||||
self.action_space = spaces.Box(
|
||||
@@ -130,35 +136,79 @@ class PendulumEnv(gym.Env):
|
||||
return np.array([np.cos(theta), np.sin(theta), thetadot], dtype=np.float32)
|
||||
|
||||
def render(self, mode="human"):
|
||||
if self.viewer is None:
|
||||
from gym.utils import pyglet_rendering
|
||||
if self.screen is None:
|
||||
pygame.init()
|
||||
self.screen = pygame.display.set_mode((self.screen_dim, self.screen_dim))
|
||||
self.surf = pygame.Surface((self.screen_dim, self.screen_dim))
|
||||
self.surf.fill((255, 255, 255))
|
||||
|
||||
self.viewer = pyglet_rendering.Viewer(500, 500)
|
||||
self.viewer.set_bounds(-2.2, 2.2, -2.2, 2.2)
|
||||
rod = pyglet_rendering.make_capsule(1, 0.2)
|
||||
rod.set_color(0.8, 0.3, 0.3)
|
||||
self.pole_transform = pyglet_rendering.Transform()
|
||||
rod.add_attr(self.pole_transform)
|
||||
self.viewer.add_geom(rod)
|
||||
axle = pyglet_rendering.make_circle(0.05)
|
||||
axle.set_color(0, 0, 0)
|
||||
self.viewer.add_geom(axle)
|
||||
fname = path.join(path.dirname(__file__), "assets/clockwise.png")
|
||||
self.img = pyglet_rendering.Image(fname, 1.0, 1.0)
|
||||
self.imgtrans = pyglet_rendering.Transform()
|
||||
self.img.add_attr(self.imgtrans)
|
||||
bound = 2.2
|
||||
scale = self.screen_dim / (bound * 2)
|
||||
offset = self.screen_dim // 2
|
||||
|
||||
self.viewer.add_onetime(self.img)
|
||||
self.pole_transform.set_rotation(self.state[0] + np.pi / 2)
|
||||
rod_length = 1 * scale
|
||||
rod_width = 0.2 * scale
|
||||
l, r, t, b = 0, rod_length, rod_width / 2, -rod_width / 2
|
||||
coords = [(l, b), (l, t), (r, t), (r, b)]
|
||||
transformed_coords = []
|
||||
for c in coords:
|
||||
c = pygame.math.Vector2(c).rotate_rad(self.state[0] + np.pi / 2)
|
||||
c = (c[0] + offset, c[1] + offset)
|
||||
transformed_coords.append(c)
|
||||
gfxdraw.aapolygon(self.surf, transformed_coords, (204, 77, 77))
|
||||
gfxdraw.filled_polygon(self.surf, transformed_coords, (204, 77, 77))
|
||||
|
||||
gfxdraw.aacircle(self.surf, offset, offset, int(rod_width / 2), (204, 77, 77))
|
||||
gfxdraw.filled_circle(
|
||||
self.surf, offset, offset, int(rod_width / 2), (204, 77, 77)
|
||||
)
|
||||
|
||||
rod_end = (rod_length, 0)
|
||||
rod_end = pygame.math.Vector2(rod_end).rotate_rad(self.state[0] + np.pi / 2)
|
||||
rod_end = (int(rod_end[0] + offset), int(rod_end[1] + offset))
|
||||
gfxdraw.aacircle(
|
||||
self.surf, rod_end[0], rod_end[1], int(rod_width / 2), (204, 77, 77)
|
||||
)
|
||||
gfxdraw.filled_circle(
|
||||
self.surf, rod_end[0], rod_end[1], int(rod_width / 2), (204, 77, 77)
|
||||
)
|
||||
|
||||
fname = path.join(path.dirname(__file__), "assets/clockwise.png")
|
||||
img = pygame.image.load(fname)
|
||||
if self.last_u is not None:
|
||||
self.imgtrans.scale = (-self.last_u / 2, np.abs(self.last_u) / 2)
|
||||
scale_img = pygame.transform.smoothscale(
|
||||
img, (scale * np.abs(self.last_u) / 2, scale * np.abs(self.last_u) / 2)
|
||||
)
|
||||
is_flip = self.last_u > 0
|
||||
scale_img = pygame.transform.flip(scale_img, is_flip, True)
|
||||
self.surf.blit(
|
||||
scale_img,
|
||||
(
|
||||
offset - scale_img.get_rect().centerx,
|
||||
offset - scale_img.get_rect().centery,
|
||||
),
|
||||
)
|
||||
|
||||
return self.viewer.render(return_rgb_array=mode == "rgb_array")
|
||||
# drawing axle
|
||||
gfxdraw.aacircle(self.surf, offset, offset, int(0.05 * scale), (0, 0, 0))
|
||||
gfxdraw.filled_circle(self.surf, offset, offset, int(0.05 * scale), (0, 0, 0))
|
||||
|
||||
self.surf = pygame.transform.flip(self.surf, False, True)
|
||||
self.screen.blit(self.surf, (0, 0))
|
||||
if mode == "human":
|
||||
pygame.display.flip()
|
||||
|
||||
if mode == "rgb_array":
|
||||
return np.transpose(
|
||||
np.array(pygame.surfarray.pixels3d(self.screen)), axes=(1, 0, 2)
|
||||
)
|
||||
else:
|
||||
return self.isopen
|
||||
|
||||
def close(self):
|
||||
if self.viewer:
|
||||
self.viewer.close()
|
||||
self.viewer = None
|
||||
if self.screen is not None:
|
||||
pygame.quit()
|
||||
self.isopen = False
|
||||
|
||||
|
||||
def angle_normalize(x):
|
||||
|
||||
@@ -1,456 +0,0 @@
|
||||
"""
|
||||
2D rendering framework
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
|
||||
if "Apple" in sys.version:
|
||||
if "DYLD_FALLBACK_LIBRARY_PATH" in os.environ:
|
||||
os.environ["DYLD_FALLBACK_LIBRARY_PATH"] += ":/usr/lib"
|
||||
# (JDS 2016/04/15): avoid bug on Anaconda 2.3.0 / Yosemite
|
||||
|
||||
from gym import error
|
||||
|
||||
try:
|
||||
import pyglet
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"""
|
||||
Cannot import pyglet.
|
||||
HINT: you can install pyglet directly via 'pip install pyglet'.
|
||||
But if you really just want to install all Gym dependencies and not have to think about it,
|
||||
'pip install -e .[all]' or 'pip install gym[all]' will do it.
|
||||
"""
|
||||
)
|
||||
|
||||
try:
|
||||
from pyglet.gl import *
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"""
|
||||
Error occurred while running `from pyglet.gl import *`
|
||||
HINT: make sure you have OpenGL installed. On Ubuntu, you can run 'apt-get install python-opengl'.
|
||||
If you're running on a server, you may need a virtual frame buffer; something like this should work:
|
||||
'xvfb-run -s \"-screen 0 1400x900x24\" python <your_script.py>'
|
||||
"""
|
||||
)
|
||||
|
||||
import math
|
||||
import numpy as np
|
||||
|
||||
RAD2DEG = 57.29577951308232
|
||||
|
||||
|
||||
def get_display(spec):
|
||||
"""Convert a display specification (such as :0) into an actual Display
|
||||
object.
|
||||
|
||||
Pyglet only supports multiple Displays on Linux.
|
||||
"""
|
||||
if spec is None:
|
||||
return pyglet.canvas.get_display()
|
||||
# returns already available pyglet_display,
|
||||
# if there is no pyglet display available then it creates one
|
||||
elif isinstance(spec, str):
|
||||
return pyglet.canvas.Display(spec)
|
||||
else:
|
||||
raise error.Error(
|
||||
f"Invalid display specification: {spec}. (Must be a string like :0 or None.)"
|
||||
)
|
||||
|
||||
|
||||
def get_window(width, height, display, **kwargs):
|
||||
"""
|
||||
Will create a pyglet window from the display specification provided.
|
||||
"""
|
||||
screen = display.get_screens() # available screens
|
||||
config = screen[0].get_best_config() # selecting the first screen
|
||||
context = config.create_context(None) # create GL context
|
||||
|
||||
return pyglet.window.Window(
|
||||
width=width,
|
||||
height=height,
|
||||
display=display,
|
||||
config=config,
|
||||
context=context,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
class Viewer:
|
||||
def __init__(self, width, height, display=None):
|
||||
display = get_display(display)
|
||||
|
||||
self.width = width
|
||||
self.height = height
|
||||
self.window = get_window(width=width, height=height, display=display)
|
||||
self.window.on_close = self.window_closed_by_user
|
||||
self.isopen = True
|
||||
self.geoms = []
|
||||
self.onetime_geoms = []
|
||||
self.transform = Transform()
|
||||
|
||||
glEnable(GL_BLEND)
|
||||
glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA)
|
||||
|
||||
def close(self):
|
||||
if self.isopen and sys.meta_path:
|
||||
# ^^^ check sys.meta_path to avoid 'ImportError: sys.meta_path is None, Python is likely shutting down'
|
||||
self.window.close()
|
||||
self.isopen = False
|
||||
|
||||
def window_closed_by_user(self):
|
||||
self.isopen = False
|
||||
|
||||
def set_bounds(self, left, right, bottom, top):
|
||||
assert right > left and top > bottom
|
||||
scalex = self.width / (right - left)
|
||||
scaley = self.height / (top - bottom)
|
||||
self.transform = Transform(
|
||||
translation=(-left * scalex, -bottom * scaley), scale=(scalex, scaley)
|
||||
)
|
||||
|
||||
def add_geom(self, geom):
|
||||
self.geoms.append(geom)
|
||||
|
||||
def add_onetime(self, geom):
|
||||
self.onetime_geoms.append(geom)
|
||||
|
||||
def render(self, return_rgb_array=False):
|
||||
glClearColor(1, 1, 1, 1)
|
||||
self.window.clear()
|
||||
self.window.switch_to()
|
||||
self.window.dispatch_events()
|
||||
self.transform.enable()
|
||||
for geom in self.geoms:
|
||||
geom.render()
|
||||
for geom in self.onetime_geoms:
|
||||
geom.render()
|
||||
self.transform.disable()
|
||||
arr = None
|
||||
if return_rgb_array:
|
||||
buffer = pyglet.image.get_buffer_manager().get_color_buffer()
|
||||
image_data = buffer.get_image_data()
|
||||
arr = np.frombuffer(image_data.get_data(), dtype=np.uint8)
|
||||
# In https://github.com/openai/gym-http-api/issues/2, we
|
||||
# discovered that someone using Xmonad on Arch was having
|
||||
# a window of size 598 x 398, though a 600 x 400 window
|
||||
# was requested. (Guess Xmonad was preserving a pixel for
|
||||
# the boundary.) So we use the buffer height/width rather
|
||||
# than the requested one.
|
||||
arr = arr.reshape(buffer.height, buffer.width, 4)
|
||||
arr = arr[::-1, :, 0:3]
|
||||
self.window.flip()
|
||||
self.onetime_geoms = []
|
||||
return arr if return_rgb_array else self.isopen
|
||||
|
||||
# Convenience
|
||||
def draw_circle(self, radius=10, res=30, filled=True, **attrs):
|
||||
geom = make_circle(radius=radius, res=res, filled=filled)
|
||||
_add_attrs(geom, attrs)
|
||||
self.add_onetime(geom)
|
||||
return geom
|
||||
|
||||
def draw_polygon(self, v, filled=True, **attrs):
|
||||
geom = make_polygon(v=v, filled=filled)
|
||||
_add_attrs(geom, attrs)
|
||||
self.add_onetime(geom)
|
||||
return geom
|
||||
|
||||
def draw_polyline(self, v, **attrs):
|
||||
geom = make_polyline(v=v)
|
||||
_add_attrs(geom, attrs)
|
||||
self.add_onetime(geom)
|
||||
return geom
|
||||
|
||||
def draw_line(self, start, end, **attrs):
|
||||
geom = Line(start, end)
|
||||
_add_attrs(geom, attrs)
|
||||
self.add_onetime(geom)
|
||||
return geom
|
||||
|
||||
def get_array(self):
|
||||
self.window.flip()
|
||||
image_data = (
|
||||
pyglet.image.get_buffer_manager().get_color_buffer().get_image_data()
|
||||
)
|
||||
self.window.flip()
|
||||
arr = np.fromstring(image_data.get_data(), dtype=np.uint8, sep="")
|
||||
arr = arr.reshape(self.height, self.width, 4)
|
||||
return arr[::-1, :, 0:3]
|
||||
|
||||
def __del__(self):
|
||||
self.close()
|
||||
|
||||
|
||||
def _add_attrs(geom, attrs):
|
||||
if "color" in attrs:
|
||||
geom.set_color(*attrs["color"])
|
||||
if "linewidth" in attrs:
|
||||
geom.set_linewidth(attrs["linewidth"])
|
||||
|
||||
|
||||
class Geom:
|
||||
def __init__(self):
|
||||
self._color = Color((0, 0, 0, 1.0))
|
||||
self.attrs = [self._color]
|
||||
|
||||
def render(self):
|
||||
for attr in reversed(self.attrs):
|
||||
attr.enable()
|
||||
self.render1()
|
||||
for attr in self.attrs:
|
||||
attr.disable()
|
||||
|
||||
def render1(self):
|
||||
raise NotImplementedError
|
||||
|
||||
def add_attr(self, attr):
|
||||
self.attrs.append(attr)
|
||||
|
||||
def set_color(self, r, g, b):
|
||||
self._color.vec4 = (r, g, b, 1)
|
||||
|
||||
|
||||
class Attr:
|
||||
def enable(self):
|
||||
raise NotImplementedError
|
||||
|
||||
def disable(self):
|
||||
pass
|
||||
|
||||
|
||||
class Transform(Attr):
|
||||
def __init__(self, translation=(0.0, 0.0), rotation=0.0, scale=(1, 1)):
|
||||
self.set_translation(*translation)
|
||||
self.set_rotation(rotation)
|
||||
self.set_scale(*scale)
|
||||
|
||||
def enable(self):
|
||||
glPushMatrix()
|
||||
glTranslatef(
|
||||
self.translation[0], self.translation[1], 0
|
||||
) # translate to GL loc ppint
|
||||
glRotatef(RAD2DEG * self.rotation, 0, 0, 1.0)
|
||||
glScalef(self.scale[0], self.scale[1], 1)
|
||||
|
||||
def disable(self):
|
||||
glPopMatrix()
|
||||
|
||||
def set_translation(self, newx, newy):
|
||||
self.translation = (float(newx), float(newy))
|
||||
|
||||
def set_rotation(self, new):
|
||||
self.rotation = float(new)
|
||||
|
||||
def set_scale(self, newx, newy):
|
||||
self.scale = (float(newx), float(newy))
|
||||
|
||||
|
||||
class Color(Attr):
|
||||
def __init__(self, vec4):
|
||||
self.vec4 = vec4
|
||||
|
||||
def enable(self):
|
||||
glColor4f(*self.vec4)
|
||||
|
||||
|
||||
class LineStyle(Attr):
|
||||
def __init__(self, style):
|
||||
self.style = style
|
||||
|
||||
def enable(self):
|
||||
glEnable(GL_LINE_STIPPLE)
|
||||
glLineStipple(1, self.style)
|
||||
|
||||
def disable(self):
|
||||
glDisable(GL_LINE_STIPPLE)
|
||||
|
||||
|
||||
class LineWidth(Attr):
|
||||
def __init__(self, stroke):
|
||||
self.stroke = stroke
|
||||
|
||||
def enable(self):
|
||||
glLineWidth(self.stroke)
|
||||
|
||||
|
||||
class Point(Geom):
|
||||
def __init__(self):
|
||||
Geom.__init__(self)
|
||||
|
||||
def render1(self):
|
||||
glBegin(GL_POINTS) # draw point
|
||||
glVertex3f(0.0, 0.0, 0.0)
|
||||
glEnd()
|
||||
|
||||
|
||||
class FilledPolygon(Geom):
|
||||
def __init__(self, v):
|
||||
Geom.__init__(self)
|
||||
self.v = v
|
||||
|
||||
def render1(self):
|
||||
if len(self.v) == 4:
|
||||
glBegin(GL_QUADS)
|
||||
elif len(self.v) > 4:
|
||||
glBegin(GL_POLYGON)
|
||||
else:
|
||||
glBegin(GL_TRIANGLES)
|
||||
for p in self.v:
|
||||
glVertex3f(p[0], p[1], 0) # draw each vertex
|
||||
glEnd()
|
||||
|
||||
|
||||
def make_circle(radius=10, res=30, filled=True):
|
||||
points = []
|
||||
for i in range(res):
|
||||
ang = 2 * math.pi * i / res
|
||||
points.append((math.cos(ang) * radius, math.sin(ang) * radius))
|
||||
if filled:
|
||||
return FilledPolygon(points)
|
||||
else:
|
||||
return PolyLine(points, True)
|
||||
|
||||
|
||||
def make_polygon(v, filled=True):
|
||||
if filled:
|
||||
return FilledPolygon(v)
|
||||
else:
|
||||
return PolyLine(v, True)
|
||||
|
||||
|
||||
def make_polyline(v):
|
||||
return PolyLine(v, False)
|
||||
|
||||
|
||||
def make_capsule(length, width):
|
||||
l, r, t, b = 0, length, width / 2, -width / 2
|
||||
box = make_polygon([(l, b), (l, t), (r, t), (r, b)])
|
||||
circ0 = make_circle(width / 2)
|
||||
circ1 = make_circle(width / 2)
|
||||
circ1.add_attr(Transform(translation=(length, 0)))
|
||||
geom = Compound([box, circ0, circ1])
|
||||
return geom
|
||||
|
||||
|
||||
class Compound(Geom):
|
||||
def __init__(self, gs):
|
||||
Geom.__init__(self)
|
||||
self.gs = gs
|
||||
for g in self.gs:
|
||||
g.attrs = [a for a in g.attrs if not isinstance(a, Color)]
|
||||
|
||||
def render1(self):
|
||||
for g in self.gs:
|
||||
g.render()
|
||||
|
||||
|
||||
class PolyLine(Geom):
|
||||
def __init__(self, v, close):
|
||||
Geom.__init__(self)
|
||||
self.v = v
|
||||
self.close = close
|
||||
self.linewidth = LineWidth(1)
|
||||
self.add_attr(self.linewidth)
|
||||
|
||||
def render1(self):
|
||||
glBegin(GL_LINE_LOOP if self.close else GL_LINE_STRIP)
|
||||
for p in self.v:
|
||||
glVertex3f(p[0], p[1], 0) # draw each vertex
|
||||
glEnd()
|
||||
|
||||
def set_linewidth(self, x):
|
||||
self.linewidth.stroke = x
|
||||
|
||||
|
||||
class Line(Geom):
|
||||
def __init__(self, start=(0.0, 0.0), end=(0.0, 0.0)):
|
||||
Geom.__init__(self)
|
||||
self.start = start
|
||||
self.end = end
|
||||
self.linewidth = LineWidth(1)
|
||||
self.add_attr(self.linewidth)
|
||||
|
||||
def render1(self):
|
||||
glBegin(GL_LINES)
|
||||
glVertex2f(*self.start)
|
||||
glVertex2f(*self.end)
|
||||
glEnd()
|
||||
|
||||
|
||||
class Image(Geom):
|
||||
def __init__(self, fname, width, height):
|
||||
Geom.__init__(self)
|
||||
self.set_color(1.0, 1.0, 1.0)
|
||||
self.width = width
|
||||
self.height = height
|
||||
img = pyglet.image.load(fname)
|
||||
self.img = img
|
||||
self.flip = False
|
||||
|
||||
def render1(self):
|
||||
self.img.blit(
|
||||
-self.width / 2, -self.height / 2, width=self.width, height=self.height
|
||||
)
|
||||
|
||||
|
||||
# ================================================================
|
||||
|
||||
|
||||
class SimpleImageViewer:
|
||||
def __init__(self, display=None, maxwidth=500):
|
||||
self.window = None
|
||||
self.isopen = False
|
||||
self.display = get_display(display)
|
||||
self.maxwidth = maxwidth
|
||||
|
||||
def imshow(self, arr):
|
||||
if self.window is None:
|
||||
height, width, _channels = arr.shape
|
||||
if width > self.maxwidth:
|
||||
scale = self.maxwidth / width
|
||||
width = int(scale * width)
|
||||
height = int(scale * height)
|
||||
self.window = get_window(
|
||||
width=width,
|
||||
height=height,
|
||||
display=self.display,
|
||||
vsync=False,
|
||||
resizable=True,
|
||||
)
|
||||
self.width = width
|
||||
self.height = height
|
||||
self.isopen = True
|
||||
|
||||
@self.window.event
|
||||
def on_resize(width, height):
|
||||
self.width = width
|
||||
self.height = height
|
||||
|
||||
@self.window.event
|
||||
def on_close():
|
||||
self.isopen = False
|
||||
|
||||
assert len(arr.shape) == 3, "You passed in an image with the wrong number shape"
|
||||
image = pyglet.image.ImageData(
|
||||
arr.shape[1], arr.shape[0], "RGB", arr.tobytes(), pitch=arr.shape[1] * -3
|
||||
)
|
||||
texture = image.get_texture()
|
||||
gl.glTexParameteri(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_MAG_FILTER, gl.GL_NEAREST)
|
||||
texture.width = self.width
|
||||
texture.height = self.height
|
||||
self.window.clear()
|
||||
self.window.switch_to()
|
||||
self.window.dispatch_events()
|
||||
texture.blit(0, 0) # draw
|
||||
self.window.flip()
|
||||
|
||||
def close(self):
|
||||
if self.isopen and sys.meta_path:
|
||||
# ^^^ check sys.meta_path to avoid 'ImportError: sys.meta_path is None, Python is likely shutting down'
|
||||
self.window.close()
|
||||
self.isopen = False
|
||||
|
||||
def __del__(self):
|
||||
self.close()
|
||||
2
setup.py
2
setup.py
@@ -13,7 +13,7 @@ extras = {
|
||||
"atari": ["ale-py~=0.7.1"],
|
||||
"accept-rom-license": ["autorom[accept-rom-license]~=0.4.2"],
|
||||
"box2d": ["box2d-py==2.3.5", "pygame==2.1.0"],
|
||||
"classic_control": ["pyglet>=1.4.0"],
|
||||
"classic_control": ["pygame==2.1.0"],
|
||||
"mujoco": ["mujoco_py>=1.50, <2.0"],
|
||||
"toy_text": ["pygame==2.1.0", "scipy>=1.4.1"],
|
||||
"other": ["lz4>=3.1.0", "opencv-python>=3.0"],
|
||||
|
||||
Reference in New Issue
Block a user