wehub-resource-sync
41 lines
1.1 KiB
Python
41 lines
1.1 KiB
Python
# __rllib-custom-gym-env-begin__
|
|
import gymnasium as gym
|
|
import numpy as np
|
|
|
|
import ray
|
|
from ray.rllib.algorithms.ppo import PPOConfig
|
|
|
|
|
|
class SimpleCorridor(gym.Env):
|
|
def __init__(self, config):
|
|
self.end_pos = config["corridor_length"]
|
|
self.cur_pos = 0.0
|
|
self.action_space = gym.spaces.Discrete(2) # right/left
|
|
self.observation_space = gym.spaces.Box(0.0, self.end_pos, shape=(1,))
|
|
|
|
def reset(self, *, seed=None, options=None):
|
|
self.cur_pos = 0.0
|
|
return np.array([self.cur_pos]), {}
|
|
|
|
def step(self, action):
|
|
if action == 0 and self.cur_pos > 0.0: # move right (towards goal)
|
|
self.cur_pos -= 1.0
|
|
elif action == 1: # move left (towards start)
|
|
self.cur_pos += 1.0
|
|
if self.cur_pos >= self.end_pos:
|
|
return np.array([0.0]), 1.0, True, True, {}
|
|
else:
|
|
return np.array([self.cur_pos]), -0.1, False, False, {}
|
|
|
|
|
|
ray.init()
|
|
|
|
config = PPOConfig().environment(SimpleCorridor, env_config={"corridor_length": 5})
|
|
algo = config.build()
|
|
|
|
for _ in range(3):
|
|
print(algo.train())
|
|
|
|
algo.stop()
|
|
# __rllib-custom-gym-env-end__
|