Python OpenAI Gym - カスタムGym環境の学習

自作の環境（カスタムGym環境）を学習させて実行してみます。

カスタムGym環境の作成

右への移動を学ぶ環境GoRightを実装します。（前回実行したもののと同じものとなります。）

エージェントが左右に移動する５マスの環境になります。

[コード]

go_right.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
import numpy as np
import gym

# 右への移動を学ぶ環境
class GoRight(gym.Env):
    # 定数定義
    GRID_SIZE = 5
    LEFT = 0
    RIGHT = 1

    # 初期化
    def __init__(self):
        super(GoRight, self).__init__()
        # グリッドのサイズ
        self.grid_size = self.GRID_SIZE
        # 初期位置の指定
        self.agent_pos = self.GRID_SIZE - 1
        # 行動空間と状態空間の定義
        self.action_space = gym.spaces.Discrete(2)
        self.observation_space = gym.spaces.Box(low=0, high=self.GRID_SIZE - 1, shape=(1,), dtype=np.float32)

    # 環境のリセット
    def reset(self):
        # 初期位置の指定
        self.agent_pos = 0
        # 初期位置をfloat32のnumpy配列に変換
        return np.array(self.agent_pos).astype(np.float32)

    # 環境の１ステップ実行
    def step(self, action):
        # 移動
        if action == self.LEFT:
            self.agent_pos -= 1
        elif action == self.RIGHT:
            self.agent_pos += 1
        self.agent_pos = np.clip(self.agent_pos, 0, self.GRID_SIZE)
        # エピソード完了の計算
        done = self.agent_pos == self.GRID_SIZE - 1
        # 報酬の計算
        reward = 1 if done else - 0.1
        return np.array(self.agent_pos).astype(np.float32), reward, done, {}

    # 環境の描画
    def render(self, mode='console', close=False):
        # エージェントはA、他は.で表現する
        print('.' * self.agent_pos, end='')
        print('A', end='')
        print('.' * (self.GRID_SIZE - 1 - self.agent_pos))

カスタムGym環境を学習させて実行

PPOで学習させて実行します。

[コード]

train_go_right.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
import gym
from go_right import GoRight

from stable_baselines.common.vec_env import DummyVecEnv
from stable_baselines import PPO2

# 環境の生成
env = GoRight()
env = DummyVecEnv([lambda: env])

# モデルの生成
model = PPO2('MlpPolicy', env, verbose=1)

# モデルの読み込み
#model = PPO2.load('go_right_model')

# モデルの学習
model.learn(total_timesteps=12800)

# モデルの保存
model.save('go_right_model')

# モデルのテスト
state = env.reset()
total_reward = 0
while True:
    # 環境の描画
    env.render()

    # モデルの推論
    action, _ = model.predict(state)

    # １ステップの実行
    state, reward, done, info = env.step(action)
    total_reward += reward

    print('')
    # エピソード完了
    if done:
        # 環境の描画
        print('total_reward:', total_reward)
        break

上記コードを実行するとコンソールに次のような表示がされます。

[実行結果]

A....

.A...

..A..

...A.

total_reward: [0.7]

学習がうまくいっているようで、一直線に右に向かって移動していることが分かります。