痛RL環境を作りたい
強化学習(Reinforcement Learning, RL)を使って研究をしているとき，
研究や仕事のモチベを上げるために痛エディタがあるなら
痛RL環境があってもいいじゃない．
と思いたったので作ってみた．
かわいい．
OpenAIGymのCartPoleのPole部分を画像に差し替える
OpenAIGymのパッケージのフォルダを開いて
gym/envs に MyEnv フォルダを作る．
中身は以下のとおり．
gym/envs/MyEnv
├ __init__.py
├ cartpole_img.py
└ assets/
└ pole_img.png
gym/envs/MyEnv/cartpole_img.py の中身
CartPoleEnv_Img クラスを作り，
既存の gym/envs/classic_control/cartpole.py の中にある CartPoleEnv クラスの
render メソッドをコピペしてきてちょっと変える．
ちなみに画像読み込み周りの処理は gym/envs/classic_control/pendulum.py からのコピペ．
gym/envs/MyEnv/cartpole_img.py
from os import path
from gym.envs.classic_control.cartpole import CartPoleEnv
class CartPoleEnv_Img(CartPoleEnv): # 既存のCartPoleEnvクラスを継承
def render(self, mode='human'):
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 * (2 * self.length)
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部分を作る元々の処理をコメントアウト
# pole = rendering.FilledPolygon([(l,b), (l,t), (r,t), (r,b)])
# pole.set_color(.8,.6,.4)
# 改変部分
fname = path.join(path.dirname(__file__), "assets/pole_pic.png") # 画像指定
scale = 200.
pole = rendering.Image(fname, scale, scale*2.) # 画像の大きさ指定
pole.set_color(1, 1, 1) # 画像の色指定
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)
self._pole_geom = pole
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])
return self.viewer.render(return_rgb_array = mode=='rgb_array')
gym/envs/MyEnv/__init__.py の中身
作ったクラスをインポート．
gym/envs/MyEnv/__init__.py
from gym.envs.MyEnv.cartpole_img import CartPoleEnv_Img
gym/envs/MyEnv/assets/pole_pic.png
pole_pic.png には好きな画像を用意すればいいが，
棒の支点にあたる部分が中心に来るような画像にしておく．
筆者はAzPainterを使って作った．
自作環境を使えるようにする
一つ上の階層 gym/envs/ にある __init__.py に，
自作環境の情報を追加しておく．
gym/envs/__init__.py
～～～省略～～～
# Classic
# ----------------------------------------
register(
id='CartPole-v0',
entry_point='gym.envs.classic_control:CartPoleEnv',
max_episode_steps=200,
reward_threshold=195.0,
)
register(
id='CartPole-v1',
entry_point='gym.envs.classic_control:CartPoleEnv',
max_episode_steps=500,
reward_threshold=475.0,
)
# ここに追加
register(
id='CartPoleImg-v0',
entry_point='gym.envs.MyEnv:CartPoleEnv_Img',
max_episode_steps=200,
reward_threshold=195.0,
)
register(
id='MountainCar-v0',
entry_point='gym.envs.classic_control:MountainCarEnv',
max_episode_steps=200,
reward_threshold=-110.0,
)
register(
id='MountainCarContinuous-v0',
entry_point='gym.envs.classic_control:Continuous_MountainCarEnv',
max_episode_steps=999,
reward_threshold=90.0,
)
register(
id='Pendulum-v0',
entry_point='gym.envs.classic_control:PendulumEnv',
max_episode_steps=200,
)
register(
id='Acrobot-v1',
entry_point='gym.envs.classic_control:AcrobotEnv',
reward_threshold=-100.0,
max_episode_steps=500,
)
～～～省略～～～
これで，env = gym.make("CartPoleImg-v0") できるようになる．
一応
版権絵を使うときは個人利用に留めましょう．
怒られても筆者は責任を負わない．