import gymnasium as gym
import stable_baselines3 as sb3
env = gym.make('[env_name]',render_mode='human')
# 定义好参数，已经学习
agent = sb3.load('[saved_model_dir]') 

observation,info = env.reset(seed=42)
terminated = false
while not terminated:
action = agent.predict(observation)
observation,reward,terminated,truncated,info = env.step(action)
if truncated:
env.reset()
env.closed

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": 30}

    def __init__(self, arg1, arg2, ...):
        super().__init__()
        # Define action and observation space
        # They must be gym.spaces objects
        # Example when using discrete actions:
        self.action_space = spaces.Discrete(N_DISCRETE_ACTIONS)
        # Example for using image as input (channel-first; channel-last also works):
        self.observation_space = spaces.Box(low=0, high=255,
                                            shape=(N_CHANNELS, HEIGHT, WIDTH), 
                                            dtype=np.uint8)

    def step(self, action):
        ...
        return observation, reward, terminated, truncated, info

    def reset(self, seed=None, options=None):
        ...
        return observation, info

    def render(self):
        ...

    def close(self):
        ...

import gymnasium
from gymnasium import spaces
import numpy as np
# Path: modelTimetable/DRL/myEnv.ipynb
# Implementing the environment
# Reproduction of the cartpole environment
# 
# Discription: 
# Create a car in a two-dimensional plane with a width of 20, and the coordinates of 
# the center point are the destination of the car to reach.
#
# State:
# The state of the car is represented by the coordinates of the center point of the car.(x,y)
# Action:
# The action of the car is represented by the speed of the car.(vx,vy)
# Reward:
# The reward is the distance between the car and the destination.
# Termination:
# The car reaches the destination.(0,0)
# truncation:
# The car is out of the screen.

'''
gymnasium is the main class that we will use to create our environment.

The gymnasium class has the following methods:
__init__(): This method is used to initialize the environment. It takes the following parameters:

step(): This method is used to take an action and return the next state, reward, and whether the episode is over. 
Physical engine
- input: action
- output: observation, reward,terminated,truncated,info

reset(): This method is used to reset the environment to its initial state.
- input: None
- output: observation

render(): This method is used to render the environment:
Image engine
- input: mode(default='human','human','rgb_array','ansi','rgb_array_list)
- output: None
eg:gymnasium.make('CartPole-v0',render_mode='human')

close(): This method is used to close the environment.
'''

class MyCar(gymnasium.Env):
    metadata = {
        'render.modes': ['human', 'rgb_array'],
        'video.frames_per_second': 2
        }
    def __init__(self):
        self.target_x = 0
        self.target_y = 0
        self.size = 10
        self.action_space = spaces.Discrete(5) # 0:stop, 1:up, 2:down, 3:left, 4:right
        self.observation_space = spaces.Box(np.array([-self.size,-self.size]), np.array([self.size,self.size]))
        self.state = None
        self.info = {}
    
    def step(self, action):
        assert self.action_space.contains(action), "%r (%s) invalid"%(action, type(action))
        # update the state by the action
        x,y = self.state
        if action == 0:
            x += 0
            y += 0
        elif action == 1:
            x += 0
            y += 1
        elif action == 2:
            x += 0
            y += -1
        elif action == 3:
            x += -1
            y += 0
        elif action == 4:
            x += 1
            y += 0
        # the next state
        self.state = np.array([x,y])
        self.state = self.state.astype(np.float32)
        reward = self._get_reward()
        terminated = self._get_terminated()
        terminated = bool(terminated)
        truncated = self._get_truncated()
        truncated = bool(truncated)
        info = {}
        return self.state, reward, terminated,truncated, info
    
    def reset(self,seed=None):
        self.state = np.ceil(np.random.rand(2)*2*self.size)-self.size
        self.state = self.state.astype(np.float32)
        self.counts = 0
        self.info = {}
        return self.state,self.info
    
    def render(self, mode='human'):
        print(self.state)
    
    def close(self):
        return super().close()

    def _get_reward(self):
        return -np.sqrt(self.state[0]**2+self.state[1]**2)
    
    def _get_terminated(self):
        x,y = self.state
        return x==self.target_x and y==self.target_y
    
    def _get_truncated(self):
        x,y = self.state
        return x<-self.size or x>self.size or y<-self.size or y>self.size

from stable_baselines3.common.env_checker import check_env
env = MyCar()
check_env(env, warn=True)