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

added rgb_array render mode for off-screen rendering

Browse Source
This commit is contained in:
Kayen 2023-11-25 23:34:24 +01:00
parent dbd3caebb3
commit deaca46d87
No known key found for this signature in database
GPG Key ID: D218E8BAD1412889
1 changed files with 20 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -8,7 +8,7 @@ from fancy_gym.envs.mujoco.air_hockey.utils import robot_to_world
from mushroom_rl.core import Environment
class AirHockeyEnv(Environment):
metadata = {"render_modes": ["human"], "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):
"""
@ -39,6 +39,18 @@ class AirHockeyEnv(Environment):
if env_mode == "tournament" and type(interpolation_order) != tuple:
interpolation_order = (interpolation_order, interpolation_order)
self.render_mode = render_mode
# Determine headless mode based on render_mode
headless = self.render_mode == 'rgb_array'
width = kwargs.pop('width', 1920)
height = kwargs.pop('height', 1080)
# Prepare viewer_params
viewer_params = kwargs.get('viewer_params', {})
viewer_params.update({'headless': headless, 'width': width, 'height': height})
kwargs['viewer_params'] = viewer_params
self.base_env = env_dict[env_mode](interpolation_order=interpolation_order, **kwargs)
self.env_name = env_mode
self.env_info = self.base_env.env_info
@ -81,9 +93,6 @@ class AirHockeyEnv(Environment):
self.env_info['constraints'] = constraint_list
self.env_info['env_name'] = self.env_name
self.render_mode = render_mode
self.render_human_active = False
super().__init__(self.base_env.info)
def step(self, action):
@ -111,15 +120,17 @@ class AirHockeyEnv(Environment):
if self.env_info['env_name'] == "tournament":
obs = np.array(np.split(obs, 2))
if self.render_human_active:
self.base_env.render()
return obs, reward, done, False, info
def render(self):
self.render_human_active = True
if self.render_mode == 'rgb_array':
return self.base_env.render(record = True)
elif self.render_mode == 'human':
self.base_env.render()
else:
raise ValueError(f"Unsupported render mode: '{self.render_mode}'")
def reset(self, seed=None, options={}):
self.base_env.seed(seed)
obs = self.base_env.reset()