"""Central configuration for the SimEngine.""" import os import numpy as np import gymnasium as gym import mujoco import numpy as np from gymnasium import spaces os.environ["MUJOCO_GL"] = "egl" # Joint names in correct order (matches real SO-101) JOINT_NAMES = [ "shoulder_pan", "shoulder_lift", "elbow_flex", "wrist_flex", "wrist_roll", "gripper", ] N_JOINTS = 6 # arm joints (excluding gripper) N_ACTIONS = 7 # 6 arm joints + 1 gripper # Joint limits in radians (used by MuJoCo actuators) JOINT_LIMITS_RAD = { "shoulder_pan": (-np.pi, np.pi), "shoulder_lift": (-np.pi / 2, np.pi / 2), "elbow_flex": (-np.pi, 0.0), "wrist_flex": (-np.pi / 2, np.pi / 2), "wrist_roll": (-np.pi, np.pi), "gripper": (0.0, 0.04), # finger separation in meters } # Home position in radians (arm straight up, out of workspace) HOME_QPOS_RAD = np.array([0.0, 0.0, -0.3, -0.3, 0.0, 0.04]) # Camera settings CAMERA_WIDTH = 640 CAMERA_HEIGHT = 480 CAMERA_NAMES = ["webcam", "arm_cam"] # Physics CONTROL_FREQ = 30 # Hz PHYSICS_TIMESTEP = 0.002 # seconds PHYSICS_SUBSTEPS = int(1.0 / CONTROL_FREQ / PHYSICS_TIMESTEP) # Episode #MAX_EPISODE_STEPS = 500 # ~17 seconds at 30 Hz # Task DEFAULT_TASK = "Pick up the box and place it in the bowl" # Paths #BASE_DIR = "/content" # жёстко задаём корень проекта в Colab #ASSETS_DIR = os.path.join(BASE_DIR, "assets") SCENE_XML = os.path.join("so101_scene.xml") # Проверяем, что сцена существует, иначе будет понятная ошибка if not os.path.exists(SCENE_XML): raise FileNotFoundError(f"Scene XML not found at {SCENE_XML}. Did you clone the repo and check assets/?") DEFAULT_DATASET_DIR = os.path.join("datasets", "sim_demos") DEFAULT_OUTPUT_DIR = os.path.join("outputs", "sim_engine") # Создаём папки, если их нет os.makedirs(DEFAULT_DATASET_DIR, exist_ok=True) os.makedirs(DEFAULT_OUTPUT_DIR, exist_ok=True) # Object positions (meters, relative to world origin at table center) BOX_DEFAULT_POS = np.array([0.15, 0.0, 0.44]) # on table surface BOWL_DEFAULT_POS = np.array([-0.15, 0.0, 0.41]) # on table surface BOWL_RADIUS = 0.05 BOWL_HEIGHT = 0.04 # Python executable (use the venv) PYTHON = "python3" class OffscreenRenderer: """Renders camera images from MuJoCo scenes.""" def __init__(self, model: mujoco.MjModel, width: int = 640, height: int = 480): self.model = model self.width = width self.height = height self._renderer = mujoco.Renderer(model, height, width) def render(self, data: mujoco.MjData, camera_name: str) -> np.ndarray: self._renderer.update_scene(data, camera=camera_name) return self._renderer.render().copy() def close(self): if hasattr(self, '_renderer'): self._renderer.close() """Gymnasium environment for SO-101 pick-and-place in MuJoCo.""" class SO101PickPlaceEnv(gym.Env): metadata = {"render_modes": ["human", "rgb_arrхay"], "render_fps": CONTROL_FREQ} def __init__(self, max_episode_steps, render_mode="rgb_array", render_images=True): super().__init__() self.render_mode = render_mode self.render_images = render_images self.max_episode_steps = max_episode_steps # Load MuJoCo model self.model = mujoco.MjModel.from_xml_path(SCENE_XML) self.model.opt.timestep = PHYSICS_TIMESTEP self.data = mujoco.MjData(self.model) # Camera renderer if self.render_images: self._renderer = OffscreenRenderer( self.model, CAMERA_WIDTH, CAMERA_HEIGHT ) else: self._renderer = None # Cache body/joint/actuator IDs self._arm_joint_ids = [ mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, name) for name in JOINT_NAMES[:5] # 5 arm joints ] self._finger_left_jnt = mujoco.mj_name2id( self.model, mujoco.mjtObj.mjOBJ_JOINT, "finger_left" ) self._box_body = mujoco.mj_name2id( self.model, mujoco.mjtObj.mjOBJ_BODY, "box" ) self._bowl_body = mujoco.mj_name2id( self.model, mujoco.mjtObj.mjOBJ_BODY, "bowl" ) self._ee_site = mujoco.mj_name2id( self.model, mujoco.mjtObj.mjOBJ_SITE, "ee_site" ) # Actuator control ranges self._ctrl_ranges = self.model.actuator_ctrlrange.copy() # Snap-grasp state (kinematic attachment when gripper closes near box) self._grasped = False self._grasp_offset = np.zeros(3) # box offset from EE when grasped self._grasp_threshold = 0.06 # EE-to-box distance to trigger grasp self._gripper_close_threshold = 0.018 # finger qpos below this = closed self._box_qpos_adr = self.model.jnt_qposadr[ mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, "box_joint") ] self._box_dof_adr = self.model.jnt_dofadr[ mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, "box_joint") ] # Remember initial state for reset mujoco.mj_forward(self.model, self.data) self._init_qpos = self.data.qpos.copy() self._init_qvel = self.data.qvel.copy() # Spaces self.action_space = spaces.Box(-1.0, 1.0, shape=(7,), dtype=np.float32) obs_spaces = {"state": spaces.Box(-np.inf, np.inf, shape=(7,), dtype=np.float32)} if self.render_images: obs_spaces["images"] = spaces.Dict({ name: spaces.Box(0, 255, shape=(CAMERA_HEIGHT, CAMERA_WIDTH, 3), dtype=np.uint8) for name in CAMERA_NAMES }) self.observation_space = spaces.Dict(obs_spaces) self._step_count = 0 self._viewer = None # ------------------------------------------------------------------ # Gym interface # ------------------------------------------------------------------ def reset(self, seed=None, options=None): super().reset(seed=seed) # Reset state self.data.qpos[:] = self._init_qpos self.data.qvel[:] = self._init_qvel self.data.ctrl[:] = 0.0 # Set arm to home position for i, jname in enumerate(JOINT_NAMES[:5]): jid = self._arm_joint_ids[i] qadr = self.model.jnt_qposadr[jid] self.data.qpos[qadr] = HOME_QPOS_RAD[i] # Open gripper fl_qadr = self.model.jnt_qposadr[self._finger_left_jnt] self.data.qpos[fl_qadr] = 0.04 # fully open # Randomize box position slightly box_qpos_adr = self.model.jnt_qposadr[ mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, "box_joint") ] self.data.qpos[box_qpos_adr:box_qpos_adr + 3] = BOX_DEFAULT_POS + self.np_random.uniform( low=[-0.03, -0.03, 0.0], high=[0.03, 0.03, 0.0] ) # Reset box orientation to identity quaternion self.data.qpos[box_qpos_adr + 3:box_qpos_adr + 7] = [1, 0, 0, 0] # Zero box velocity box_dof_adr = self.model.jnt_dofadr[ mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, "box_joint") ] self.data.qvel[box_dof_adr:box_dof_adr + 6] = 0.0 # Reset grasp state self._grasped = False self._grasp_offset = np.zeros(3) # Set actuator controls to match initial positions for i in range(5): jid = self._arm_joint_ids[i] qadr = self.model.jnt_qposadr[jid] self.data.ctrl[i] = self.data.qpos[qadr] self.data.ctrl[5] = 0.04 # left finger open self.data.ctrl[6] = 0.04 # right finger open # Settle physics for _ in range(50): mujoco.mj_step(self.model, self.data) self._step_count = 0 obs = self._get_obs() return obs, self._get_info() def step(self, action): action = np.clip(action, -1.0, 1.0).astype(np.float32) # Scale action from [-1,1] to actuator control ranges # First 5: arm joints, 6th: gripper (controls both fingers) for i in range(5): lo, hi = self._ctrl_ranges[i] self.data.ctrl[i] = lo + (action[i] + 1.0) * 0.5 * (hi - lo) # Gripper: action[5] in [-1,1] -> finger separation [0.005, 0.04] # Min 0.005 prevents fingers from squeezing box out gripper_val = 0.005 + (action[5] + 1.0) * 0.5 * 0.035 self.data.ctrl[5] = gripper_val # left finger self.data.ctrl[6] = gripper_val # right finger (coupled but set explicitly) # action[6] is the second gripper dim (ignored if coupled, or can be used for asymmetric grip) # For simplicity, we use action[5] for both fingers # Update snap-grasp (based on gripper action command) self._update_grasp(action[5]) # Step physics for _ in range(PHYSICS_SUBSTEPS): mujoco.mj_step(self.model, self.data) self._step_count += 1 obs = self._get_obs() info = self._get_info() reward = self._compute_reward(info) terminated = info["success"] truncated = self._step_count >= self.max_episode_steps return obs, reward, terminated, truncated, info def render(self): if self.render_mode == "rgb_array" and self._renderer is not None: return self._renderer.render(self.data, "webcam") elif self.render_mode == "human": self._render_human() return None def close(self): if self._renderer is not None: self._renderer.close() if self._viewer is not None: self._viewer.close() # ------------------------------------------------------------------ # Internal # ------------------------------------------------------------------ def _update_grasp(self, gripper_action): ee_pos = self.data.site_xpos[self._ee_site].copy() box_pos = self.data.xpos[self._box_body].copy() dist = np.linalg.norm(ee_pos - box_pos) fl_qadr = self.model.jnt_qposadr[self._finger_left_jnt] finger_qpos = self.data.qpos[fl_qadr] gripper_closing = gripper_action < -0.5 # action commands close gripper_opening = gripper_action > 0.5 # action commands open if not self._grasped and gripper_closing and finger_qpos < self._gripper_close_threshold and dist < self._grasp_threshold: self._grasp_offset = box_pos - ee_pos self._grasped = True elif self._grasped and gripper_opening: self._grasped = False if self._grasped: new_box_pos = ee_pos + self._grasp_offset self.data.qpos[self._box_qpos_adr:self._box_qpos_adr + 3] = new_box_pos self.data.qvel[self._box_dof_adr:self._box_dof_adr + 6] = 0.0 def _get_obs(self): # Joint positions in degrees joint_pos = np.zeros(5, dtype=np.float32) for i, jid in enumerate(self._arm_joint_ids): qadr = self.model.jnt_qposadr[jid] joint_pos[i] = np.rad2deg(self.data.qpos[qadr]) # Gripper opening: 0 (closed) to 100 (open) fl_qadr = self.model.jnt_qposadr[self._finger_left_jnt] gripper_pos = self.data.qpos[fl_qadr] gripper_normalized = np.float32(np.clip(gripper_pos / 0.04 * 100.0, 0, 100)) state = np.append(joint_pos, gripper_normalized).astype(np.float32) obs = {"state": state} if self.render_images and self._renderer is not None: obs["images"] = { name: self._renderer.render(self.data, name) for name in CAMERA_NAMES } return obs def _get_info(self): box_pos = self.data.xpos[self._box_body].copy() bowl_pos = self.data.xpos[self._bowl_body].copy() ee_pos = self.data.site_xpos[self._ee_site].copy() # Check if box is in the bowl horizontal_dist = np.linalg.norm(box_pos[:2] - bowl_pos[:2]) box_above_bowl_base = box_pos[2] > bowl_pos[2] - 0.01 box_below_bowl_rim = box_pos[2] < bowl_pos[2] + BOWL_HEIGHT + 0.02 success = horizontal_dist < BOWL_RADIUS and box_above_bowl_base and box_below_bowl_rim return { "success": bool(success), "box_pos": box_pos, "bowl_pos": bowl_pos, "ee_pos": ee_pos, "ee_to_box_dist": float(np.linalg.norm(ee_pos - box_pos)), "box_to_bowl_dist": float(np.linalg.norm(box_pos[:2] - bowl_pos[:2])), } def _compute_reward(self, info): reward = 0.0 # Shaping: encourage approaching the box reward -= 0.1 * info["ee_to_box_dist"] # Grasp bonus: box is lifted above table if info["box_pos"][2] > 0.47: reward += 1.0 # Shaping: encourage moving box toward bowl reward -= 0.1 * info["box_to_bowl_dist"] # Placement bonus if info["success"]: reward += 5.0 return reward def _render_human(self): """Open an interactive MuJoCo viewer.""" if self._viewer is None: import mujoco.viewer self._viewer = mujoco.viewer.launch_passive(self.model, self.data) else: self._viewer.sync() # ------------------------------------------------------------------ # Utilities # ------------------------------------------------------------------ def get_ee_pos(self): """Get end-effector position in world frame.""" return self.data.site_xpos[self._ee_site].copy() def get_box_pos(self): """Get box position in world frame.""" return self.data.xpos[self._box_body].copy() def get_bowl_pos(self): """Get bowl position in world frame.""" return self.data.xpos[self._bowl_body].copy() @property def arm_joint_ids(self): return self._arm_joint_ids @property def ee_site_id(self): return self._ee_site