An environment for an Aloha robot in simulation.
| 7 | |
| 8 | |
| 9 | class AlohaSimEnvironment(_environment.Environment): |
| 10 | """An environment for an Aloha robot in simulation.""" |
| 11 | |
| 12 | def __init__(self, task: str, obs_type: str = "pixels_agent_pos", seed: int = 0) -> None: |
| 13 | np.random.seed(seed) |
| 14 | self._rng = np.random.default_rng(seed) |
| 15 | |
| 16 | self._gym = gymnasium.make(task, obs_type=obs_type) |
| 17 | |
| 18 | self._last_obs = None |
| 19 | self._done = True |
| 20 | self._episode_reward = 0.0 |
| 21 | |
| 22 | @override |
| 23 | def reset(self) -> None: |
| 24 | gym_obs, _ = self._gym.reset(seed=int(self._rng.integers(2**32 - 1))) |
| 25 | self._last_obs = self._convert_observation(gym_obs) # type: ignore |
| 26 | self._done = False |
| 27 | self._episode_reward = 0.0 |
| 28 | |
| 29 | @override |
| 30 | def is_episode_complete(self) -> bool: |
| 31 | return self._done |
| 32 | |
| 33 | @override |
| 34 | def get_observation(self) -> dict: |
| 35 | if self._last_obs is None: |
| 36 | raise RuntimeError("Observation is not set. Call reset() first.") |
| 37 | |
| 38 | return self._last_obs # type: ignore |
| 39 | |
| 40 | @override |
| 41 | def apply_action(self, action: dict) -> None: |
| 42 | gym_obs, reward, terminated, truncated, info = self._gym.step(action["actions"]) |
| 43 | self._last_obs = self._convert_observation(gym_obs) # type: ignore |
| 44 | self._done = terminated or truncated |
| 45 | self._episode_reward = max(self._episode_reward, reward) |
| 46 | |
| 47 | def _convert_observation(self, gym_obs: dict) -> dict: |
| 48 | img = gym_obs["pixels"]["top"] |
| 49 | img = image_tools.convert_to_uint8(image_tools.resize_with_pad(img, 224, 224)) |
| 50 | # Convert axis order from [H, W, C] --> [C, H, W] |
| 51 | img = np.transpose(img, (2, 0, 1)) |
| 52 | |
| 53 | return { |
| 54 | "state": gym_obs["agent_pos"], |
| 55 | "images": {"cam_high": img}, |
| 56 | } |
nothing calls this directly
no outgoing calls
no test coverage detected