File size: 2,097 Bytes
079c32c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 |
from typing import TYPE_CHECKING, Optional
from numpy import ndarray
if TYPE_CHECKING:
from ding.envs import BaseEnv, BaseEnvManager
def render_env(env, render_mode: Optional[str] = 'rgb_array') -> "ndarray":
"""
Overview:
Render the environment's current frame.
Arguments:
- env (:obj:`gym.Env`): DI-engine env instance.
- render_mode (:obj:`str`): Render mode.
Returns:
- frame (:obj:`numpy.ndarray`): [H * W * C]
"""
if hasattr(env, 'sim'):
# mujoco: mujoco frame is unside-down by default
return env.sim.render(camera_name='track', height=128, width=128)[::-1]
else:
# other envs
return env.render(mode=render_mode)
def render(env: "BaseEnv", render_mode: Optional[str] = 'rgb_array') -> "ndarray":
"""
Overview:
Render the environment's current frame.
Arguments:
- env (:obj:`BaseEnv`): DI-engine env instance.
- render_mode (:obj:`str`): Render mode.
Returns:
- frame (:obj:`numpy.ndarray`): [H * W * C]
"""
gym_env = env._env
return render_env(gym_env, render_mode=render_mode)
def get_env_fps(env) -> "int":
"""
Overview:
Get the environment's fps.
Arguments:
- env (:obj:`gym.Env`): DI-engine env instance.
Returns:
- fps (:obj:`int`).
"""
if hasattr(env, 'model'):
# mujoco
fps = 1 / env.model.opt.timestep
elif hasattr(env, 'env') and 'video.frames_per_second' in env.env.metadata.keys():
# classic control
fps = env.env.metadata['video.frames_per_second']
else:
# atari and other envs
fps = 30
return fps
def fps(env_manager: "BaseEnvManager") -> "int":
"""
Overview:
Render the environment's fps.
Arguments:
- env (:obj:`BaseEnvManager`): DI-engine env manager instance.
Returns:
- fps (:obj:`int`).
"""
try:
# env_ref is a ding gym environment
gym_env = env_manager.env_ref._env
return get_env_fps(gym_env)
except:
return 30
|