31 lines
933 B
Python
31 lines
933 B
Python
import gymnasium as gym
|
|
import numpy as np
|
|
from gymnasium import spaces
|
|
|
|
class CustomEnv(gym.Env):
|
|
"""Custom Environment that follows gym interface."""
|
|
|
|
metadata = {"render_modes": ["human"], "render_fps": 1}
|
|
|
|
def __init__(self, s_dim=0):
|
|
super().__init__()
|
|
self.S_DIM = s_dim
|
|
self.action_space = spaces.Box(low=-1, high=1, shape=(3,), dtype=np.float32)
|
|
self.observation_space = spaces.Box(
|
|
low=-1, high=1, shape=(s_dim,), dtype=np.float32
|
|
)
|
|
|
|
def step(self, action):
|
|
return np.zeros(self.S_DIM, dtype=np.float32), float(0), False, False, {}
|
|
|
|
def change_s_dim(self, s_dim=0):
|
|
self.S_DIM = s_dim
|
|
self.observation_space = spaces.Box(
|
|
low=-1, high=1, shape=(s_dim,), dtype=np.float32
|
|
)
|
|
|
|
def reset(self, seed=None):
|
|
return np.zeros(self.S_DIM, dtype=np.float32), {}
|
|
|
|
def close(self):
|
|
self.__del__() |