gym 自定义游戏环境

需要继承gym.Env类，实现reset(self)、step(self, action)函数

import gym
from typing import List, Optional, Tuple, Any


class ToyEnv(gym.Env):
    def __init__(self):
        super(ToyEnv, self).__init__()
        self.observation_space = gym.spaces.Discrete(n=5)   # 定义状态空间为discrete类型 范围0~4
        self.action_space = gym.spaces.Discrete(n=3)        # 定义动作空间为discrete类型 范围0~2
        self.step_index = 0

    def reset(self):    # 初始化状态，计数器置零
        self.step_index = 0
        return self.step_index

    def step(self, action):
        is_done = self.step_index == 10     # 计数器到10，游戏结束
        if is_done:     # 游戏结束
            # 返回(s,r,done,info)
            return self.step_index % self.observation_space.n, \
                   0.0, is_done, {}

        # 游戏未结束 返回(s_,r,done,info)
        self.step_index += 1
        return self.step_index % self.observation_space.n, \
               float(action), False, {}

# 使用
env = ToyEnv()
s = env.reset()
print("env.reset() -> %s" % s)
s = env.step(1)