Update sim/base_sim.py

sim/base_sim.py

@@ -2,7 +2,7 @@ import argparse
 import pathlib
 from pathlib import Path
 import threading
-from threading import Lock, Thread
+from threading import Thread
 from typing import Dict
 
 import mujoco
@@ -18,7 +18,7 @@ from unitree_sdk2py.core.channel import ChannelFactoryInitialize
 import yaml
 import os
 from .image_publish_utils import ImagePublishProcess
-from .metric_utils import check_contact, check_height
+from .metric_utils import check_contact
 from .sim_utils import get_subtree_body_names
 from .unitree_sdk2py_bridge import ElasticBand, UnitreeSdk2Bridge
 
@@ -60,8 +60,6 @@ class DefaultEnv:
         if len(camera_configs) > 0:
             print(f"✓ DefaultEnv initialized with {len(camera_configs)} camera(s): {list(camera_configs.keys())}")
 
-        # Thread safety lock
-        self.reward_lock = Lock()
 
         # Unitree bridge will be initialized by the simulator
         self.unitree_bridge = None
@@ -71,7 +69,6 @@ class DefaultEnv:
 
         # Initialize scene (defined in subclasses)
         self.init_scene()
-        self.last_reward = 0
 
         # Setup offscreen rendering if needed
         self.offscreen = offscreen
@@ -439,16 +436,6 @@ class DefaultEnv:
             else:
                 self.viewer.cam.type = mujoco.mjtCamera.mjCAMERA_TRACKING
 
-    def update_reward(self):
-        """Calculate reward. Should be implemented by subclasses."""
-        with self.reward_lock:
-            self.last_reward = 0
-
-    def get_reward(self):
-        """Thread-safe way to get the last calculated reward."""
-        with self.reward_lock:
-            return self.last_reward
-
     def set_unitree_bridge(self, unitree_bridge):
         """Set the unitree bridge from the simulator"""
         self.unitree_bridge = unitree_bridge
@@ -509,114 +496,6 @@ class DefaultEnv:
         mujoco.mj_resetData(self.mj_model, self.mj_data)
 
 
-class CubeEnv(DefaultEnv):
-    """Environment with a cube object for pick and place tasks"""
-
-    def __init__(
-        self,
-        config: Dict[str, any],
-        onscreen: bool = False,
-        offscreen: bool = False,
-    ):
-        # Override the robot scene
-        config = config.copy()  # Create a copy to avoid modifying the original
-        config["ROBOT_SCENE"] = "gr00t_wbc/control/robot_model/model_data/g1/pnp_cube_43dof.xml"
-        super().__init__(config, "cube", {}, onscreen, offscreen)
-
-    def update_reward(self):
-        """Calculate reward based on gripper contact with cube and cube height"""
-        right_hand_body = [
-            "right_hand_thumb_2_link",
-            "right_hand_middle_1_link",
-            "right_hand_index_1_link",
-        ]
-        gripper_cube_contact = check_contact(
-            self.mj_model, self.mj_data, right_hand_body, "cube_body"
-        )
-        cube_lifted = check_height(self.mj_model, self.mj_data, "cube", 0.85, 2.0)
-
-        with self.reward_lock:
-            self.last_reward = gripper_cube_contact & cube_lifted
-
-
-class BoxEnv(DefaultEnv):
-    """Environment with a box object for manipulation tasks"""
-
-    def __init__(
-        self,
-        config: Dict[str, any],
-        onscreen: bool = False,
-        offscreen: bool = False,
-    ):
-        # Override the robot scene
-        config = config.copy()  # Create a copy to avoid modifying the original
-        config["ROBOT_SCENE"] = "gr00t_wbc/control/robot_model/model_data/g1/lift_box_43dof.xml"
-        super().__init__(config, "box", {}, onscreen, offscreen)
-
-    def reward(self):
-        """Calculate reward based on gripper contact with cube and cube height"""
-        left_hand_body = [
-            "left_hand_thumb_2_link",
-            "left_hand_middle_1_link",
-            "left_hand_index_1_link",
-        ]
-        right_hand_body = [
-            "right_hand_thumb_2_link",
-            "right_hand_middle_1_link",
-            "right_hand_index_1_link",
-        ]
-        gripper_box_contact = check_contact(self.mj_model, self.mj_data, left_hand_body, "box_body")
-        gripper_box_contact &= check_contact(
-            self.mj_model, self.mj_data, right_hand_body, "box_body"
-        )
-        box_lifted = check_height(self.mj_model, self.mj_data, "box", 0.92, 2.0)
-
-        print("gripper_box_contact: ", gripper_box_contact, "box_lifted: ", box_lifted)
-
-        with self.reward_lock:
-            self.last_reward = gripper_box_contact & box_lifted
-            return self.last_reward
-
-
-class BottleEnv(DefaultEnv):
-    """Environment with a cylinder object for manipulation tasks"""
-
-    def __init__(
-        self,
-        config: Dict[str, any],
-        onscreen: bool = False,
-        offscreen: bool = False,
-    ):
-        # Override the robot scene
-        config = config.copy()  # Create a copy to avoid modifying the original
-        config["ROBOT_SCENE"] = "gr00t_wbc/control/robot_model/model_data/g1/pnp_bottle_43dof.xml"
-        camera_configs = {
-            "egoview": {
-                "height": 400,
-                "width": 400,
-            },
-        }
-        super().__init__(
-            config, "cylinder", camera_configs, onscreen, offscreen
-        )
-
-        self.bottle_body = self.mj_model.body("bottle_body")
-        self.bottle_geom = self.mj_model.geom("bottle")
-
-        if self.viewer is not None:
-            self.viewer.cam.type = mujoco.mjtCamera.mjCAMERA_FIXED
-            self.viewer.cam.fixedcamid = self.mj_model.camera("egoview").id
-
-    def update_reward(self):
-        """Calculate reward based on gripper contact with cylinder and cylinder height"""
-        pass
-
-    def get_privileged_obs(self):
-        obs_pos = self.mj_data.xpos[self.bottle_body.id]
-        obs_quat = self.mj_data.xquat[self.bottle_body.id]
-        return {"bottle_pos": obs_pos, "bottle_quat": obs_quat}
-
-
 class BaseSimulator:
     """Base simulator class that handles initialization and running of simulations"""
 
@@ -641,21 +520,11 @@ class BaseSimulator:
 
         # Set update frequencies
         self.sim_dt = self.config["SIMULATE_DT"]
-        self.reward_dt = self.config.get("REWARD_DT", 0.02)
         self.image_dt = self.config.get("IMAGE_DT", 0.033333)
         self.viewer_dt = self.config.get("VIEWER_DT", 0.02)
 
-        # Create the appropriate environment based on name
-        if env_name == "default":
-            self.sim_env = DefaultEnv(config, env_name, **kwargs)
-        elif env_name == "pnp_cube":
-            self.sim_env = CubeEnv(config, **kwargs)
-        elif env_name == "lift_box":
-            self.sim_env = BoxEnv(config, **kwargs)
-        elif env_name == "pnp_bottle":
-            self.sim_env = BottleEnv(config, **kwargs)
-        else:
-            raise ValueError(f"Invalid environment name: {env_name}")
+        # Create the environment
+        self.sim_env = DefaultEnv(config, env_name, **kwargs)
 
         # Initialize the DDS communication layer - should be safe to call multiple times
 
@@ -722,10 +591,6 @@ class BaseSimulator:
                 if sim_cnt % int(self.viewer_dt / self.sim_dt) == 0:
                     self.sim_env.update_viewer()
 
-                # Calculate reward at reward rate
-                if sim_cnt % int(self.reward_dt / self.sim_dt) == 0:
-                    self.sim_env.update_reward()
-
                 # Update render caches at image rate
                 if sim_cnt % int(self.image_dt / self.sim_dt) == 0:
                     self.sim_env.update_render_caches()