dodox/fancy_gym
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

incorporated human render_mode gym spec + optimized keyword arguments

Browse Source
This commit is contained in:
Kayen 2023-11-26 21:49:52 +01:00
parent deaca46d87
commit 0b4e729a49
No known key found for this signature in database
GPG Key ID: D218E8BAD1412889
1 changed files with 6 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
fancy_gym/envs/mujoco/air_hockey/air_hockey_env_wrapper.py
View File

@ -10,7 +10,7 @@ from mushroom_rl.core import Environment
class AirHockeyEnv(Environment): class AirHockeyEnv(Environment):
metadata = {"render_modes": ["human", "rgb_array"], "render_fps": 50} metadata = {"render_modes": ["human", "rgb_array"], "render_fps": 50}
def __init__(self, env_mode=None, interpolation_order=3, render_mode=None, **kwargs): def __init__(self, env_mode=None, interpolation_order=3, render_mode=None, width=1920, height=1080, **kwargs):
""" """
Environment Constructor Environment Constructor
@ -40,11 +40,10 @@ class AirHockeyEnv(Environment):
interpolation_order = (interpolation_order, interpolation_order) interpolation_order = (interpolation_order, interpolation_order)
self.render_mode = render_mode self.render_mode = render_mode
self.render_human_active = False
# Determine headless mode based on render_mode # Determine headless mode based on render_mode
headless = self.render_mode == 'rgb_array' headless = self.render_mode == 'rgb_array'
width = kwargs.pop('width', 1920)
height = kwargs.pop('height', 1080)
# Prepare viewer_params # Prepare viewer_params
viewer_params = kwargs.get('viewer_params', {}) viewer_params = kwargs.get('viewer_params', {})
@ -121,13 +120,16 @@ class AirHockeyEnv(Environment):
if self.env_info['env_name'] == "tournament": if self.env_info['env_name'] == "tournament":
obs = np.array(np.split(obs, 2)) obs = np.array(np.split(obs, 2))
if self.render_human_active:
self.base_env.render()
return obs, reward, done, False, info return obs, reward, done, False, info
def render(self): def render(self):
if self.render_mode == 'rgb_array': if self.render_mode == 'rgb_array':
return self.base_env.render(record = True) return self.base_env.render(record = True)
elif self.render_mode == 'human': elif self.render_mode == 'human':
self.base_env.render() self.render_human_active = True
else: else:
raise ValueError(f"Unsupported render mode: '{self.render_mode}'") raise ValueError(f"Unsupported render mode: '{self.render_mode}'")