Permalink
29 lines (25 sloc) 769 Bytes
import numpy as np
import gym
class Discrete(gym.Space):
"""
{0,1,...,n-1}
Example usage:
self.observation_space = spaces.Discrete(2)
"""
def __init__(self, n):
self.n = n
gym.Space.__init__(self, (), np.int64)
def sample(self):
return gym.spaces.np_random.randint(self.n)
def contains(self, x):
if isinstance(x, int):
as_int = x
elif isinstance(x, (np.generic, np.ndarray)) and (x.dtype.kind in np.typecodes['AllInteger'] and x.shape == ()):
as_int = int(x)
else:
return False
return as_int >= 0 and as_int < self.n
def __repr__(self):
return "Discrete(%d)" % self.n
def __eq__(self, other):
return self.n == other.n