import gymnasium as gym


class MetaworldAddRenderingToInfoWrapper(gym.Wrapper):
    def __init__(self, env: gym.Env, key: str = "rendering"):
        super().__init__(env)
        self.key = key

    def step(self, action):
        observation, reward, terminated, truncated, info = self.env.step(action)
        # render rgb image with default camera
        info[self.key] = self.env.unwrapped.mujoco_renderer.render(render_mode="rgb_array")  # type: ignore
        return observation, reward, terminated, truncated, info


from .build import build

__all__ = ["build"]
