bowl and microwave env

examples/run_hang_lifebuoy_teleop.py

@@ -0,0 +1,88 @@
+#!/usr/bin/env python3
+"""
+Teleop script for Hang Lifebuoy environment.
+"""
+
+import genesis as gs
+import torch
+from gs_agent.wrappers.teleop_wrapper import KeyboardWrapper
+from gs_env.sim.envs.config.registry import EnvArgsRegistry
+from gs_env.sim.envs.manipulation.hang_lifebuoy_env import HangLifebuoyEnv
+
+
+def main() -> None:
+    """Run teleop for hang lifebuoy task."""
+    print("Initializing Hang Lifebuoy Teleop System...")
+
+    # Initialize Genesis
+    gs.init(
+        seed=0,
+        precision="32",
+        logging_level="info",
+        backend=gs.cpu,
+    )
+
+    # Create teleop wrapper first (without environment)
+    print("Creating teleop wrapper...")
+    teleop_wrapper = KeyboardWrapper(
+        env=None,
+        device=torch.device("cpu"),
+        movement_speed=0.01,  # Position movement speed
+        rotation_speed=0.05,  # Rotation speed
+        trajectory_filename_prefix="hang_lifebuoy_",
+    )
+
+    # Start teleop wrapper (keyboard listener) FIRST, before creating Genesis scene
+    teleop_wrapper.start()
+
+    # Create task environment AFTER teleop wrapper is running
+    env = HangLifebuoyEnv(
+        args=EnvArgsRegistry["hang_lifebuoy_default"],
+        device=torch.device("cpu"),
+    )
+    teleop_wrapper.set_environment(env)
+
+    print("\n" + "=" * 50)
+    print("Hang Lifebuoy TELEOP SYSTEM READY")
+    print("=" * 50)
+    print("📝 TRAJECTORY RECORDING INSTRUCTIONS:")
+    print("   1. Press 'r' to start recording (anytime)")
+    print("   2. Move robot with arrow keys, n/m, space")
+    print("   3. Press 'r' again to stop recording and save")
+    print("   💡 Recording works from any state now!")
+    print("=" * 50)
+
+    # Run the main control loop in the main thread (Genesis viewer requires this)
+    try:
+        step_count = 0
+        while teleop_wrapper.running:
+            # Step the teleop wrapper (this processes input and steps environment)
+            teleop_wrapper.step(torch.tensor([]))
+            step_count += 1
+
+            # Check for quit command
+            if (
+                hasattr(teleop_wrapper, "last_command")
+                and teleop_wrapper.last_command
+                and hasattr(teleop_wrapper.last_command, "quit_teleop")
+                and teleop_wrapper.last_command.quit_teleop
+            ):
+                print("Quit command received, exiting...")
+                break
+
+            # Safety check - exit after 1 hour of running
+            if step_count > 180000:  # 1 hour at 50Hz
+                print("Maximum runtime reached, exiting...")
+                break
+
+    except KeyboardInterrupt:
+        print("\n👋 Teleop interrupted by user")
+
+    finally:
+        # Cleanup
+        teleop_wrapper.stop()
+        print("✅ Teleop session ended")
+
+
+if __name__ == "__main__":
+    main()

examples/run_pick_cube_teleop.py

@@ -35,7 +35,7 @@ def main() -> None:
         seed=0,
         precision="32",
         logging_level="info",
-        backend=gs.cpu,  # type: ignore
+        backend=gs.cpu,
     )
 
     try:
@@ -49,7 +49,7 @@ def main() -> None:
         )
 
         # Start teleop wrapper (keyboard listener) FIRST, before creating Genesis scene
-        teleop_wrapper.start()  # type: ignore
+        teleop_wrapper.start()
 
         # Create task environment AFTER teleop wrapper is running
         env = create_gs_env()

examples/run_put_bowl_teleop.py

@@ -0,0 +1,88 @@
+#!/usr/bin/env python3
+"""
+Teleop script for Put Bowl Inside Microwave environment.
+"""
+
+import genesis as gs
+import torch
+from gs_agent.wrappers.teleop_wrapper import KeyboardWrapper
+from gs_env.sim.envs.config.registry import EnvArgsRegistry
+from gs_env.sim.envs.manipulation.put_bowl_inside_microwave_env import PutBowlInsideMicrowaveEnv
+
+
+def main() -> None:
+    """Run teleop for put bowl inside microwave task."""
+    print("Initializing Put Bowl Inside Microwave Teleop System...")
+
+    # Initialize Genesis
+    gs.init(
+        seed=0,
+        precision="32",
+        logging_level="info",
+        backend=gs.cpu,  # type: ignore
+    )
+
+    # Create teleop wrapper first (without environment)
+    print("Creating teleop wrapper...")
+    teleop_wrapper = KeyboardWrapper(
+        env=None,
+        device=torch.device("cpu"),
+        movement_speed=0.01,  # Position movement speed
+        rotation_speed=0.05,  # Rotation speed
+        trajectory_filename_prefix="put_bowl_microwave_",
+    )
+
+    # Start teleop wrapper (keyboard listener) FIRST, before creating Genesis scene
+    teleop_wrapper.start()  # type: ignore
+
+    # Create task environment AFTER teleop wrapper is running
+    env = PutBowlInsideMicrowaveEnv(
+        args=EnvArgsRegistry["put_bowl_inside_microwave_default"],
+        device=torch.device("cpu"),
+    )
+    teleop_wrapper.set_environment(env)
+
+    print("\n" + "=" * 50)
+    print("Put Bowl Inside Microwave TELEOP SYSTEM READY")
+    print("=" * 50)
+    print("📝 TRAJECTORY RECORDING INSTRUCTIONS:")
+    print("   1. Press 'r' to start recording (anytime)")
+    print("   2. Move robot with arrow keys, n/m, space")
+    print("   3. Press 'r' again to stop recording and save")
+    print("   💡 Recording works from any state now!")
+    print("=" * 50)
+
+    # Run the main control loop in the main thread (Genesis viewer requires this)
+    try:
+        step_count = 0
+        while teleop_wrapper.running:
+            # Step the teleop wrapper (this processes input and steps environment)
+            teleop_wrapper.step(torch.tensor([]))
+            step_count += 1
+
+            # Check for quit command
+            if (
+                hasattr(teleop_wrapper, "last_command")
+                and teleop_wrapper.last_command
+                and hasattr(teleop_wrapper.last_command, "quit_teleop")
+                and teleop_wrapper.last_command.quit_teleop
+            ):
+                print("Quit command received, exiting...")
+                break
+
+            # Safety check - exit after 1 hour of running
+            if step_count > 180000:  # 1 hour at 50Hz
+                print("Maximum runtime reached, exiting...")
+                break
+
+    except KeyboardInterrupt:
+        print("\n👋 Teleop interrupted by user")
+
+    finally:
+        # Cleanup
+        teleop_wrapper.stop()
+        print("✅ Teleop session ended")
+
+
+if __name__ == "__main__":
+    main()

examples/run_sweep_table_teleop.py

@@ -0,0 +1,88 @@
+#!/usr/bin/env python3
+"""
+Teleop script for Sweep Table environment.
+"""
+
+import genesis as gs
+import torch
+from gs_agent.wrappers.teleop_wrapper import KeyboardWrapper
+from gs_env.sim.envs.config.registry import EnvArgsRegistry
+from gs_env.sim.envs.manipulation.sweep_table_env import SweepTableEnv
+
+
+def main() -> None:
+    """Run teleop for sweep table task."""
+    print("Initializing Sweep Table Teleop System...")
+
+    # Initialize Genesis
+    gs.init(
+        seed=0,
+        precision="32",
+        logging_level="info",
+        backend=gs.cpu,
+    )
+
+    # Create teleop wrapper first (without environment)
+    print("Creating teleop wrapper...")
+    teleop_wrapper = KeyboardWrapper(
+        env=None,
+        device=torch.device("cpu"),
+        movement_speed=0.01,  # Position movement speed
+        rotation_speed=0.05,  # Rotation speed
+        trajectory_filename_prefix="sweep_table_",
+    )
+
+    # Start teleop wrapper (keyboard listener) FIRST, before creating Genesis scene
+    teleop_wrapper.start()
+
+    # Create task environment AFTER teleop wrapper is running
+    env = SweepTableEnv(
+        args=EnvArgsRegistry["sweep_table_default"],
+        device=torch.device("cpu"),
+    )
+    teleop_wrapper.set_environment(env)
+
+    print("\n" + "=" * 50)
+    print("Sweep Table TELEOP SYSTEM READY")
+    print("=" * 50)
+    print("📝 TRAJECTORY RECORDING INSTRUCTIONS:")
+    print("   1. Press 'r' to start recording (anytime)")
+    print("   2. Move robot with arrow keys, n/m, space")
+    print("   3. Press 'r' again to stop recording and save")
+    print("   💡 Recording works from any state now!")
+    print("=" * 50)
+
+    # Run the main control loop in the main thread (Genesis viewer requires this)
+    try:
+        step_count = 0
+        while teleop_wrapper.running:
+            # Step the teleop wrapper (this processes input and steps environment)
+            teleop_wrapper.step(torch.tensor([]))
+            step_count += 1
+
+            # Check for quit command
+            if (
+                hasattr(teleop_wrapper, "last_command")
+                and teleop_wrapper.last_command
+                and hasattr(teleop_wrapper.last_command, "quit_teleop")
+                and teleop_wrapper.last_command.quit_teleop
+            ):
+                print("Quit command received, exiting...")
+                break
+
+            # Safety check - exit after 1 hour of running
+            if step_count > 180000:  # 1 hour at 50Hz
+                print("Maximum runtime reached, exiting...")
+                break
+
+    except KeyboardInterrupt:
+        print("\n👋 Teleop interrupted by user")
+
+    finally:
+        # Cleanup
+        teleop_wrapper.stop()
+        print("✅ Teleop session ended")
+
+
+if __name__ == "__main__":
+    main()

src/agent/gs_agent/bases/env_wrapper.py

@@ -14,15 +14,15 @@ def __init__(self, env: Any, device: torch.device) -> None:
         self.device: Final[torch.device] = device
 
     @abstractmethod
-    def reset(self) -> tuple[torch.Tensor, dict[str, Any]]: ...
+    def reset(self) -> tuple[dict[str, Any], dict[str, Any]]: ...
 
     @abstractmethod
     def step(
         self, action: torch.Tensor
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, dict[str, Any]]: ...
+    ) -> tuple[dict[str, Any], torch.Tensor, torch.Tensor, torch.Tensor, dict[str, Any]]: ...
 
     @abstractmethod
-    def get_observations(self) -> torch.Tensor: ...
+    def get_observations(self) -> dict[str, Any]: ...
 
     @property
     @abstractmethod

src/agent/gs_agent/wrappers/gs_env_wrapper.py

@@ -22,14 +22,14 @@ def __init__(
     # ---------------------------
     # BatchEnvWrapper API (batch)
     # ---------------------------
-    def reset(self) -> tuple[torch.Tensor, dict[str, Any]]:
+    def reset(self) -> tuple[dict[str, Any], dict[str, Any]]:
         self.env.reset()
         self._curr_obs = self.env.get_observations()
         return self._curr_obs, self.env.get_extra_infos()
 
     def step(
         self, action: torch.Tensor
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, dict[str, Any]]:
+    ) -> tuple[dict[str, Any], torch.Tensor, torch.Tensor, torch.Tensor, dict[str, Any]]:
         # apply action
         self.env.apply_action(action)
         # get observations
@@ -55,7 +55,7 @@ def step(
         extra_infos["reward_terms"] = reward_terms
         return next_obs, reward, terminated, truncated, extra_infos
 
-    def get_observations(self) -> torch.Tensor:
+    def get_observations(self) -> dict[str, Any]:
         return self.env.get_observations()
 
     @property

src/agent/gs_agent/wrappers/gym_env_wrapper.py

@@ -20,14 +20,14 @@ def __init__(self, env: gym.Env[Any, Any], device: torch.device = _DEFAULT_DEVIC
     # ---------------------------
     # BatchEnvWrapper API (batch)
     # ---------------------------
-    def reset(self) -> tuple[torch.Tensor, dict[str, Any]]:
+    def reset(self) -> tuple[dict[str, Any], dict[str, Any]]:
         obs, info = self.env.reset()
         self._curr_obs = torch.tensor(obs, device=self.device).unsqueeze(0)
-        return self._curr_obs, info
+        return {"obs": self._curr_obs}, info
 
     def step(
         self, action: torch.Tensor
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, dict[str, Any]]:
+    ) -> tuple[dict[str, Any], torch.Tensor, torch.Tensor, torch.Tensor, dict[str, Any]]:
         gym_actions = action.clone().cpu().numpy().squeeze(0)  # shape: [action_dim]
 
         obs, reward, terminated, truncated, info = self.env.step(gym_actions)
@@ -42,10 +42,10 @@ def step(
         reward_batch = torch.as_tensor([[float(reward)]], device=self.device)
         terminated_batch = torch.as_tensor([[terminated]], device=self.device)
         truncated_batch = torch.as_tensor([[truncated]], device=self.device)
-        return self._curr_obs, reward_batch, terminated_batch, truncated_batch, info
+        return {"obs": self._curr_obs}, reward_batch, terminated_batch, truncated_batch, info
 
-    def get_observations(self) -> torch.Tensor:
-        return self._curr_obs
+    def get_observations(self) -> dict[str, Any]:
+        return {"obs": self._curr_obs}
 
     @property
     def action_dim(self) -> int: