diff --git a/panda_gym/__init__.py b/panda_gym/__init__.py index cf7865c7..aa5f7ec6 100644 --- a/panda_gym/__init__.py +++ b/panda_gym/__init__.py @@ -1,3 +1,5 @@ +"""Gymnasium registration for Panda and safety-aware Panda environments.""" + import os from gymnasium.envs.registration import register @@ -6,124 +8,36 @@ __version__ = file_handler.read().strip() ENV_IDS = [] - -for task in ["Reach", "Slide", "Push", "PickAndPlace", "Stack", "Flip",\ - "ReachSafe", "PushSafe", "SlideSafe", "PickAndPlaceSafe", "StackSafe"\ - "Stack3", "StackPyramid", "BuildL" ]: - for reward_type in ["sparse", "dense"]: - for control_type in ["ee", "joints"]: +TASKS = ( + "Reach", + "Slide", + "Push", + "PickAndPlace", + "Stack", + "Flip", + "ReachSafe", + "PushSafe", + "SlideSafe", + "PickAndPlaceSafe", + "StackSafe", + "BuildL", +) + +for task in TASKS: + for reward_type in ("sparse", "dense"): + for control_type in ("ee", "joints"): reward_suffix = "Dense" if reward_type == "dense" else "" control_suffix = "Joints" if control_type == "joints" else "" env_id = f"Panda{task}{control_suffix}{reward_suffix}-v3" - register( id=env_id, entry_point=f"panda_gym.envs:Panda{task}Env", kwargs={"reward_type": reward_type, "control_type": control_type}, - max_episode_steps=100 if task == "Stack" else 50, + max_episode_steps=100 + if task in {"Stack", "StackSafe", "BuildL"} + else 50, ) - ENV_IDS.append(env_id) - - # register( - # id="PandaReachSafe{}{}-v2".format(control_suffix, reward_suffix), - # entry_point="panda_gym.envs:PandaReachSafeEnv", - # kwargs=kwargs, - # max_episode_steps=50, - # ) - - - # register( - # id="PandaPush{}{}-v2".format(control_suffix, reward_suffix), - # entry_point="panda_gym.envs:PandaPushEnv", - # kwargs=kwargs, - # max_episode_steps=50, - # ) - - # register( - # id="PandaPushSafe{}{}-v2".format(control_suffix, reward_suffix), - # entry_point="panda_gym.envs:PandaPushSafeEnv", - # kwargs=kwargs, - # max_episode_steps=50, - # ) - - # register( - # id="PandaSlide{}{}-v2".format(control_suffix, reward_suffix), - # entry_point="panda_gym.envs:PandaSlideEnv", - # kwargs=kwargs, - # max_episode_steps=50, - # ) - - # register( - # id="PandaSlideSafe{}{}-v2".format(control_suffix, reward_suffix), - # entry_point="panda_gym.envs:PandaSlideSafeEnv", - # kwargs=kwargs, - # max_episode_steps=50, - # ) - - # register( - # id="PandaPickAndPlace{}{}-v2".format(control_suffix, reward_suffix), - # entry_point="panda_gym.envs:PandaPickAndPlaceEnv", - # kwargs=kwargs, - # max_episode_steps=50, - # ) - - # register( - # id="PandaPickAndPlaceSafe{}{}-v2".format(control_suffix, reward_suffix), - # entry_point="panda_gym.envs:PandaPickAndPlaceSafeEnv", - # kwargs=kwargs, - # max_episode_steps=50, - # ) - - # register( - # id="PandaPickAndPlacePlatform{}{}-v2".format(control_suffix, reward_suffix), - # entry_point="panda_gym.envs:PandaPickAndPlacePlatformEnv", - # kwargs=kwargs, - # max_episode_steps=50, - # ) - - - # register( - # id="PandaStack{}{}-v2".format(control_suffix, reward_suffix), - # entry_point="panda_gym.envs:PandaStackEnv", - # kwargs=kwargs, - # max_episode_steps=100, - # ) - - # register( - # id="PandaStackSafe{}{}-v2".format(control_suffix, reward_suffix), - # entry_point="panda_gym.envs:PandaStackSafeEnv", - # kwargs=kwargs, - # max_episode_steps=100, - # ) - - - # register( - # id="PandaStack3{}{}-v2".format(control_suffix, reward_suffix), - # entry_point="panda_gym.envs:PandaStack3Env", - # kwargs=kwargs, - # max_episode_steps=100, - # ) - # register( - # id="PandaStackPyramid{}{}-v2".format(control_suffix, reward_suffix), - # entry_point="panda_gym.envs:PandaStackPyramidEnv", - # kwargs=kwargs, - # max_episode_steps=100, - # ) - - # register( - # id="PandaBuildL{}{}-v2".format(control_suffix, reward_suffix), - # entry_point="panda_gym.envs:PandaBuildLEnv", - # kwargs=kwargs, - # max_episode_steps=100, - # ) - - - # register( - # id="PandaFlip{}{}-v2".format(control_suffix, reward_suffix), - # entry_point="panda_gym.envs:PandaFlipEnv", - # kwargs=kwargs, - # max_episode_steps=50, - # ) +__all__ = ["ENV_IDS", "TASKS", "__version__"] diff --git a/panda_gym/envs/__init__.py b/panda_gym/envs/__init__.py index 78fd880d..49ab9321 100644 --- a/panda_gym/envs/__init__.py +++ b/panda_gym/envs/__init__.py @@ -1,3 +1,5 @@ +"""Environment exports used by Gymnasium entry points.""" + from panda_gym.envs.panda_tasks import ( PandaFlipEnv, PandaPickAndPlaceEnv, @@ -5,27 +7,27 @@ PandaReachEnv, PandaSlideEnv, PandaStackEnv, - - #new environments - PandaStack3Env, - PandaStackPyramidEnv, - PandaPickAndPlacePlatformEnv, - - #safe environments +) +from panda_gym.envs.panda_tasks_safe import ( + PandaPickAndPlaceSafeEnv, PandaPushSafeEnv, PandaReachSafeEnv, PandaSlideSafeEnv, - PandaPickAndPlaceSafeEnv, - PandaStackSafeEnv + PandaStackSafeEnv, ) +from panda_gym.envs.panda_tasks_multi import PandaBuildLEnv -# from panda_gym.envs.panda_tasks.panda_pick_and_place_platform import PandaPickAndPlacePlatformEnv -# from panda_gym.envs.panda_tasks.panda_stack_pyramid import PandaStackPyramidEnv -#from panda_gym.envs.panda_tasks.panda_stack_3 import PandaStack3Env -#from panda_gym.envs.panda_tasks.panda_stack_pyramid import PandaStackPyramidEnv -#from panda_gym.envs.panda_tasks.panda_build_L import PandaBuildLEnv -# from panda_gym.envs.panda_tasks.panda_push_safe import PandaPushSafeEnv -# from panda_gym.envs.panda_tasks.panda_reach_safe import PandaReachSafeEnv -# from panda_gym.envs.panda_tasks.panda_slide_safe import PandaSlideSafeEnv -# from panda_gym.envs.panda_tasks.panda_pick_and_place_safe import PandaPickAndPlaceSafeEnv -# from panda_gym.envs.panda_tasks.panda_stack_safe import PandaStackSafeEnv +__all__ = [ + "PandaFlipEnv", + "PandaPickAndPlaceEnv", + "PandaPushEnv", + "PandaReachEnv", + "PandaSlideEnv", + "PandaStackEnv", + "PandaPickAndPlaceSafeEnv", + "PandaPushSafeEnv", + "PandaReachSafeEnv", + "PandaSlideSafeEnv", + "PandaStackSafeEnv", + "PandaBuildLEnv", +] diff --git a/panda_gym/envs/core_safe.py b/panda_gym/envs/core_safe.py index 9f229f44..2972c42a 100644 --- a/panda_gym/envs/core_safe.py +++ b/panda_gym/envs/core_safe.py @@ -1,336 +1,29 @@ -from abc import ABC, abstractmethod -from typing import Any, Dict, Optional, Tuple, Union +"""Gymnasium-compatible core classes for safety-aware Panda tasks.""" -import gym -import gym.spaces -import gym.utils.seeding -import gym_robotics -import numpy as np - -from panda_gym.pybullet import PyBullet - - -class PyBulletRobot(ABC): - """Base class for robot env. - - Args: - sim (PyBullet): Simulation instance. - body_name (str): The name of the robot within the simulation. - file_name (str): Path of the urdf file. - base_position (np.ndarray): Position of the base of the robot as (x, y, z). - """ - - def __init__( - self, - sim: PyBullet, - body_name: str, - file_name: str, - base_position: np.ndarray, - action_space: gym.spaces.Space, - joint_indices: np.ndarray, - joint_forces: np.ndarray, - ) -> None: - self.sim = sim - self.body_name = body_name - with self.sim.no_rendering(): - self._load_robot(file_name, base_position) - self.setup() - self.action_space = action_space - self.joint_indices = joint_indices - self.joint_forces = joint_forces - - def _load_robot(self, file_name: str, base_position: np.ndarray) -> None: - """Load the robot. - - Args: - file_name (str): The URDF file name of the robot. - base_position (np.ndarray): The position of the robot, as (x, y, z). - """ - self.sim.loadURDF( - body_name=self.body_name, - fileName=file_name, - basePosition=base_position, - useFixedBase=True, - ) - - def setup(self) -> None: - """Called after robot loading.""" - pass - - @abstractmethod - def set_action(self, action: np.ndarray) -> None: - """Set the action. Must be called just before sim.step(). - - Args: - action (np.ndarray): The action. - """ - - @abstractmethod - def get_obs(self) -> np.ndarray: - """Return the observation associated to the robot. - - Returns: - np.ndarray: The observation. - """ - - @abstractmethod - def reset(self) -> None: - """Reset the robot and return the observation.""" - - def get_link_position(self, link: int) -> np.ndarray: - """Returns the position of a link as (x, y, z) - - Args: - link (int): The link index. - - Returns: - np.ndarray: Position as (x, y, z) - """ - return self.sim.get_link_position(self.body_name, link) - - def get_link_velocity(self, link: int) -> np.ndarray: - """Returns the velocity of a link as (vx, vy, vz) - - Args: - link (int): The link index. - - Returns: - np.ndarray: Velocity as (vx, vy, vz) - """ - return self.sim.get_link_velocity(self.body_name, link) - - def get_joint_angle(self, joint: int) -> float: - """Returns the angle of a joint - - Args: - joint (int): The joint index. - - Returns: - float: Joint angle - """ - return self.sim.get_joint_angle(self.body_name, joint) - - def get_joint_velocity(self, joint: int) -> float: - """Returns the velocity of a joint as (wx, wy, wz) - - Args: - joint (int): The joint index. - - Returns: - np.ndarray: Joint velocity as (wx, wy, wz) - """ - return self.sim.get_joint_velocity(self.body_name, joint) - - def control_joints(self, target_angles: np.ndarray) -> None: - """Control the joints of the robot. - - Args: - target_angles (np.ndarray): The target angles. The length of the array must equal to the number of joints. - """ - self.sim.control_joints( - body=self.body_name, - joints=self.joint_indices, - target_angles=target_angles, - forces=self.joint_forces, - ) - - def set_joint_angles(self, angles: np.ndarray) -> None: - """Set the joint position of a body. Can induce collisions. +from typing import Any, Dict, Tuple - Args: - angles (list): Joint angles. - """ - self.sim.set_joint_angles(self.body_name, joints=self.joint_indices, angles=angles) - - def inverse_kinematics(self, link: int, position: np.ndarray, orientation: np.ndarray) -> np.ndarray: - """Compute the inverse kinematics and return the new joint values. - - Args: - link (int): The link. - position (x, y, z): Desired position of the link. - orientation (x, y, z, w): Desired orientation of the link. - - Returns: - List of joint values. - """ - inverse_kinematics = self.sim.inverse_kinematics(self.body_name, link=link, position=position, orientation=orientation) - return inverse_kinematics - - -class Task(ABC): - """Base class for tasks. - Args: - sim (PyBullet): Simulation instance. - """ - - def __init__(self, sim: PyBullet) -> None: - self.sim = sim - self.goal = None - - @abstractmethod - def reset(self) -> None: - """Reset the task: sample a new goal.""" - - @abstractmethod - def get_obs(self) -> np.ndarray: - """Return the observation associated to the task.""" - - @abstractmethod - def get_achieved_goal(self) -> np.ndarray: - """Return the achieved goal.""" - - def get_goal(self) -> np.ndarray: - """Return the current goal.""" - if self.goal is None: - raise RuntimeError("No goal yet, call reset() first") - else: - return self.goal.copy() - - - @abstractmethod - def is_success( - self, achieved_goal: np.ndarray, desired_goal: np.ndarray, info: Dict[str, Any] = {} - ) -> Union[np.ndarray, float]: - """Returns whether the achieved goal match the desired goal.""" - - @abstractmethod - def compute_reward( - self, achieved_goal: np.ndarray, desired_goal: np.ndarray, info: Dict[str, Any] = {} - ) -> Union[np.ndarray, float]: - """Compute reward associated to the achieved and the desired goal.""" - - @abstractmethod - def compute_cost( - self, achieved_goal: np.ndarray, desired_goal: np.ndarray, info: Dict[str, Any] = {} - ) -> Union[np.ndarray, float]: - """Compute reward associated to the achieved and the desired goal.""" - - -class RobotTaskEnv(gym_robotics.GoalEnv): - """Robotic task goal env, as the junction of a task and a robot. - - Args: - robot (PyBulletRobot): The robot. - task (Task): The task. - """ - - metadata = {"render.modes": ["human", "rgb_array"]} - - def __init__(self, robot: PyBulletRobot, task: Task) -> None: - assert robot.sim == task.sim, "The robot and the task must belong to the same simulation." - self.sim = robot.sim - self.robot = robot - self.task = task - obs = self.reset() # required for init; seed can be changed later - observation_shape = obs.shape - - # achieved_goal_shape = obs["achieved_goal"].shape - # desired_goal_shape = obs["achieved_goal"].shape - # self.observation_space = gym.spaces.Dict( - # dict( - # observation=gym.spaces.Box(-10.0, 10.0, shape=observation_shape, dtype=np.float32), - # desired_goal=gym.spaces.Box(-10.0, 10.0, shape=achieved_goal_shape, dtype=np.float32), - # achieved_goal=gym.spaces.Box(-10.0, 10.0, shape=desired_goal_shape, dtype=np.float32), - # ) - # ) - - self.observation_space = gym.spaces.Box(-10.0, 10.0, shape=observation_shape, dtype=np.float32) - - self.action_space = self.robot.action_space - self.compute_reward = self.task.compute_reward - self._saved_goal = dict() - - def _get_obs(self) -> Dict[str, np.ndarray]: - robot_obs = self.robot.get_obs() # robot state - task_obs = self.task.get_obs() # object position, velococity, unsafe state locations etc... - observation = np.concatenate([robot_obs, task_obs]) - # achieved_goal = self.task.get_achieved_goal() - # return np.concatenate([observation, achieved_goal, self.task.get_goal()]) - return observation - # np.concatenate([observation, achieved_goal, self.task.get_goal()]) - return { - "observation": np.concatenate([observation, achieved_goal, self.task.get_goal()]), - # "observation": observation, - # "achieved_goal": achieved_goal, - # "desired_goal": self.task.get_goal(), - } - - def reset(self, seed: Optional[int] = None) -> Dict[str, np.ndarray]: - self.task.np_random, seed = gym.utils.seeding.np_random(seed) - with self.sim.no_rendering(): - self.robot.reset() - self.task.reset() - return self._get_obs() +import numpy as np - def save_state(self) -> int: - state_id = self.sim.save_state() - self._saved_goal[state_id] = self.task.goal - return state_id +from panda_gym.envs.core import PyBulletRobot, RobotTaskEnv as BaseRobotTaskEnv, Task - def restore_state(self, state_id: int) -> None: - self.sim.restore_state(state_id) - self.task.goal = self._saved_goal[state_id] - def remove_state(self, state_id: int) -> None: - self._saved_goal.pop(state_id) - self.sim.remove_state(state_id) +class RobotTaskEnv(BaseRobotTaskEnv): + """Robot-task environment that exposes the task safety cost in ``info``.""" - def step(self, action: np.ndarray) -> Tuple[Dict[str, np.ndarray], float, bool, Dict[str, Any]]: + def step( + self, action: np.ndarray + ) -> Tuple[Dict[str, np.ndarray], float, bool, bool, Dict[str, Any]]: self.robot.set_action(action) self.sim.step() - obs = self._get_obs() - achieved_goal = self.task.get_achieved_goal() - get_goal = self.task.get_goal() - # done = False if self.task.is_success(achieved_goal, get_goal)==np.array([0]) else True - done = False - cost_value = self.task.compute_cost() - info = {"is_success": self.task.is_success(achieved_goal, get_goal), "cost" : cost_value} - reward = self.task.compute_reward(achieved_goal, get_goal, info) - assert isinstance(reward, float) # needed for pytype cheking - return obs, reward, done, info - - def close(self) -> None: - self.sim.close() - - def render( - self, - mode: str, - width: int = 720, - height: int = 480, - target_position: Optional[np.ndarray] = None, - distance: float = 1.4, - yaw: float = 45, - pitch: float = -30, - roll: float = 0, - ) -> Optional[np.ndarray]: - """Render. - - If mode is "human", make the rendering real-time. All other arguments are - unused. If mode is "rgb_array", return an RGB array of the scene. - - Args: - mode (str): "human" of "rgb_array". If "human", this method waits for the time necessary to have - a realistic temporal rendering and all other args are ignored. Else, return an RGB array. - width (int, optional): Image width. Defaults to 720. - height (int, optional): Image height. Defaults to 480. - target_position (np.ndarray, optional): Camera targetting this postion, as (x, y, z). - Defaults to [0., 0., 0.]. - distance (float, optional): Distance of the camera. Defaults to 1.4. - yaw (float, optional): Yaw of the camera. Defaults to 45. - pitch (float, optional): Pitch of the camera. Defaults to -30. - roll (int, optional): Rool of the camera. Defaults to 0. - - Returns: - RGB np.ndarray or None: An RGB array if mode is 'rgb_array', else None. - """ - target_position = target_position if target_position is not None else np.zeros(3) - return self.sim.render( - mode, - width=width, - height=height, - target_position=target_position, - distance=distance, - yaw=yaw, - pitch=pitch, - roll=roll, - ) + observation = self._get_obs() + achieved_goal = observation["achieved_goal"] + desired_goal = self.task.get_goal() + terminated = bool(self.task.is_success(achieved_goal, desired_goal)) + truncated = False + cost = float(self.task.compute_cost()) + info = {"is_success": terminated, "cost": cost} + reward = float(self.task.compute_reward(achieved_goal, desired_goal, info)) + return observation, reward, terminated, truncated, info + + +__all__ = ["PyBulletRobot", "Task", "RobotTaskEnv"] diff --git a/panda_gym/envs/panda_tasks_multi.py b/panda_gym/envs/panda_tasks_multi.py new file mode 100644 index 00000000..01bd568f --- /dev/null +++ b/panda_gym/envs/panda_tasks_multi.py @@ -0,0 +1,51 @@ +"""Gymnasium constructors for multi-object Panda tasks.""" + +from typing import Optional + +import numpy as np + +from panda_gym.envs.core import RobotTaskEnv +from panda_gym.envs.robots.panda import Panda +from panda_gym.envs.tasks.build_L_ import BuildL +from panda_gym.pybullet import PyBullet + + +class PandaBuildLEnv(RobotTaskEnv): + """Arrange four colored cubes into the translucent L-shaped targets.""" + + def __init__( + self, + render_mode: str = "rgb_array", + reward_type: str = "sparse", + control_type: str = "ee", + renderer: str = "Tiny", + render_width: int = 720, + render_height: int = 480, + render_target_position: Optional[np.ndarray] = None, + render_distance: float = 0.9, + render_yaw: float = 45, + render_pitch: float = -30, + render_roll: float = 0, + ) -> None: + sim = PyBullet(render_mode=render_mode, renderer=renderer) + robot = Panda( + sim, + block_gripper=False, + base_position=np.array([-0.6, 0.0, 0.0]), + control_type=control_type, + ) + task = BuildL(sim, reward_type=reward_type) + super().__init__( + robot, + task, + render_width=render_width, + render_height=render_height, + render_target_position=render_target_position, + render_distance=render_distance, + render_yaw=render_yaw, + render_pitch=render_pitch, + render_roll=render_roll, + ) + + +__all__ = ["PandaBuildLEnv"] diff --git a/panda_gym/envs/panda_tasks_safe.py b/panda_gym/envs/panda_tasks_safe.py new file mode 100644 index 00000000..89178af5 --- /dev/null +++ b/panda_gym/envs/panda_tasks_safe.py @@ -0,0 +1,88 @@ +"""Gymnasium constructors for the safety-aware Panda task variants.""" + +from typing import Optional, Type + +import numpy as np + +from panda_gym.envs.core_safe import RobotTaskEnv, Task +from panda_gym.envs.robots.panda import Panda +from panda_gym.envs.tasks.pick_and_place_safe import PickAndPlaceSafe +from panda_gym.envs.tasks.push_safe import Push +from panda_gym.envs.tasks.reach_safe import Reach +from panda_gym.envs.tasks.slide_safe import Slide +from panda_gym.envs.tasks.stack_safe import Stack +from panda_gym.pybullet import PyBullet + + +class _SafePandaEnv(RobotTaskEnv): + task_type: Type[Task] + block_gripper: bool = True + + def __init__( + self, + render_mode: str = "rgb_array", + reward_type: str = "sparse", + control_type: str = "ee", + renderer: str = "Tiny", + render_width: int = 720, + render_height: int = 480, + render_target_position: Optional[np.ndarray] = None, + render_distance: float = 1.4, + render_yaw: float = 45, + render_pitch: float = -30, + render_roll: float = 0, + ) -> None: + sim = PyBullet(render_mode=render_mode, renderer=renderer) + robot = Panda( + sim, + block_gripper=self.block_gripper, + base_position=np.array([-0.6, 0.0, 0.0]), + control_type=control_type, + ) + task = self.task_type( + sim, + reward_type=reward_type, + get_ee_position=robot.get_ee_position, + ) + super().__init__( + robot, + task, + render_width=render_width, + render_height=render_height, + render_target_position=render_target_position, + render_distance=render_distance, + render_yaw=render_yaw, + render_pitch=render_pitch, + render_roll=render_roll, + ) + + +class PandaReachSafeEnv(_SafePandaEnv): + task_type = Reach + + +class PandaPushSafeEnv(_SafePandaEnv): + task_type = Push + + +class PandaSlideSafeEnv(_SafePandaEnv): + task_type = Slide + + +class PandaPickAndPlaceSafeEnv(_SafePandaEnv): + task_type = PickAndPlaceSafe + block_gripper = False + + +class PandaStackSafeEnv(_SafePandaEnv): + task_type = Stack + block_gripper = False + + +__all__ = [ + "PandaReachSafeEnv", + "PandaPushSafeEnv", + "PandaSlideSafeEnv", + "PandaPickAndPlaceSafeEnv", + "PandaStackSafeEnv", +] diff --git a/panda_gym/envs/tasks/build_L_.py b/panda_gym/envs/tasks/build_L_.py index 8dd3c6a4..117e636c 100644 --- a/panda_gym/envs/tasks/build_L_.py +++ b/panda_gym/envs/tasks/build_L_.py @@ -2,7 +2,7 @@ import numpy as np -from panda_gym.envs.core_multi_task import Task +from panda_gym.envs.core import Task from panda_gym.pybullet import PyBullet from panda_gym.utils import distance diff --git a/test_safe_envs/test_gymnasium_api.py b/test_safe_envs/test_gymnasium_api.py new file mode 100644 index 00000000..29cc3b7e --- /dev/null +++ b/test_safe_envs/test_gymnasium_api.py @@ -0,0 +1,58 @@ +import gymnasium as gym +import numpy as np +import pytest + +import panda_gym + + +SAFE_ENV_IDS = [ + "PandaReachSafe-v3", + "PandaPushSafe-v3", + "PandaSlideSafe-v3", + "PandaPickAndPlaceSafe-v3", + "PandaStackSafe-v3", +] + + +def test_registry_contains_separate_safe_environment_ids() -> None: + assert set(SAFE_ENV_IDS).issubset(panda_gym.ENV_IDS) + assert not any("StackSafeStack3" in env_id for env_id in panda_gym.ENV_IDS) + assert "PandaBuildL-v3" in panda_gym.ENV_IDS + + +def test_build_l_environment_reset_step_and_render_headless() -> None: + env = gym.make("PandaBuildL-v3", render_mode="rgb_array", renderer="Tiny") + try: + observation, info = env.reset(seed=7) + assert env.observation_space.contains(observation) + assert "is_success" in info + action = np.zeros(env.action_space.shape, dtype=env.action_space.dtype) + observation, reward, terminated, truncated, info = env.step(action) + assert env.observation_space.contains(observation) + assert isinstance(reward, float) + assert isinstance(terminated, bool) + assert isinstance(truncated, bool) + frame = env.render() + assert frame.shape == (480, 720, 3) + finally: + env.close() + + +@pytest.mark.parametrize("env_id", SAFE_ENV_IDS) +def test_safe_environment_reset_and_step_headless(env_id: str) -> None: + env = gym.make(env_id, render_mode="rgb_array", renderer="Tiny") + try: + observation, info = env.reset(seed=7) + assert set(observation) == {"observation", "achieved_goal", "desired_goal"} + assert env.observation_space.contains(observation) + assert "is_success" in info + + action = np.zeros(env.action_space.shape, dtype=env.action_space.dtype) + observation, reward, terminated, truncated, info = env.step(action) + assert env.observation_space.contains(observation) + assert isinstance(reward, float) + assert isinstance(terminated, bool) + assert isinstance(truncated, bool) + assert isinstance(info["cost"], float) + finally: + env.close()