mirror of
https://github.com/vale981/ray
synced 2025-03-08 11:31:40 -05:00
42991d723f
Cleans up of the rllib/examples folder by moving all example Envs into rllibexamples/env (so they can be used by other scripts and tests as well).
44 lines
1.3 KiB
Python
44 lines
1.3 KiB
Python
import atexit
|
|
import gym
|
|
from gym.spaces import Discrete
|
|
import os
|
|
import subprocess
|
|
import time
|
|
|
|
|
|
class EnvWithSubprocess(gym.Env):
|
|
"""Our env that spawns a subprocess."""
|
|
|
|
# Dummy command to run as a subprocess with a unique name
|
|
UNIQUE_CMD = "sleep {}".format(str(time.time()))
|
|
|
|
def __init__(self, config):
|
|
self.UNIQUE_FILE_0 = config["tmp_file1"]
|
|
self.UNIQUE_FILE_1 = config["tmp_file2"]
|
|
self.UNIQUE_FILE_2 = config["tmp_file3"]
|
|
self.UNIQUE_FILE_3 = config["tmp_file4"]
|
|
|
|
self.action_space = Discrete(2)
|
|
self.observation_space = Discrete(2)
|
|
# Subprocess that should be cleaned up
|
|
self.subproc = subprocess.Popen(
|
|
self.UNIQUE_CMD.split(" "), shell=False)
|
|
self.config = config
|
|
# Exit handler should be called
|
|
atexit.register(lambda: self.subproc.kill())
|
|
if config.worker_index == 0:
|
|
atexit.register(lambda: os.unlink(self.UNIQUE_FILE_0))
|
|
else:
|
|
atexit.register(lambda: os.unlink(self.UNIQUE_FILE_1))
|
|
|
|
def close(self):
|
|
if self.config.worker_index == 0:
|
|
os.unlink(self.UNIQUE_FILE_2)
|
|
else:
|
|
os.unlink(self.UNIQUE_FILE_3)
|
|
|
|
def reset(self):
|
|
return 0
|
|
|
|
def step(self, action):
|
|
return 0, 0, True, {}
|