RL-Project/Fetch-Reinforcement_learning_Project

1# <-- this must come first, before any mujoco / gym imports2# import os3# os.environ["MUJOCO_GL"] = "egl"4 5import numpy as np6import gymnasium as gym7import gymnasium_robotics8import mujoco9 10class CustomFetchWrapper(gym.Wrapper):11    def __init__(self, env, block_xy=None, goal_xyz=None, object=True):12        super().__init__(env)13        self.u = env.unwrapped  # MujocoFetchPickAndPlaceEnv14        # stash your fixed coords (or None to randomize)15        self.default_block_xy = (np.array(block_xy, dtype=float)16                                 if block_xy is not None else None)17        self.default_goal_xyz = (np.array(goal_xyz, dtype=float)18                                 if goal_xyz is not None else None)19        self.object = object20 21    def reset(self, *args, **kwargs):22        # 1) do the normal reset — gets you a random goal in obs23        obs, info = super().reset(*args, **kwargs)24        u     = self.unwrapped25        model = u.model26        data  = u.data27        utils = u._utils28        rng   = u.np_random29 30        # 2) reset the robot slides to your home pose31        for name, val in zip(32            ["robot0:slide0","robot0:slide1","robot0:slide2"],33            [0.405,       0.48,        0.0],34        ):35            utils.set_joint_qpos(model, data, name, val)36 37        # pull out the actual goal so we can avoid it38        goal_pos = obs["desired_goal"][:2].copy()39 40        if (self.object==True):41            # 3) pick block position42            if self.default_block_xy is None:43                home_xy  = u.initial_gripper_xpos[:2]44                obj_range = u.obj_range45                min_dist  = u.distance_threshold46 47                while True:48                    offset = rng.uniform(-obj_range, obj_range, size=2)49                    # 3a) must be outside the “too-close to gripper” zone50                    if np.linalg.norm(offset) < min_dist:51                        continue52                    candidate_xy = home_xy + offset53                    # 3b) must be outside the “too-close to goal” zone54                    if np.linalg.norm(candidate_xy - goal_pos) < min_dist:55                        continue56                    # if we get here, both checks passed57                    break58 59                block_xy = candidate_xy60                61            else:62                block_xy = self.default_block_xy63 64            # place the block65            blk_qpos = utils.get_joint_qpos(model, data, "object0:joint")66            blk_qpos[0:2] = block_xy67            blk_qpos[2]    = 0.42  # table height68            utils.set_joint_qpos(model, data, "object0:joint", blk_qpos)69 70        # 4) pick goal position71        if self.default_goal_xyz is not None:72            new_goal = self.default_goal_xyz73 74            # override the goal both in the env and in the MuJoCo site75            u.goal = new_goal76            sid = mujoco.mj_name2id(model,77                                    mujoco.mjtObj.mjOBJ_SITE,78                                    "target0")79            data.site_xpos[sid] = new_goal80 81        # 5) forward‐kinematics + fresh obs82        u._mujoco.mj_forward(model, data)83        obs = u._get_obs()84 85        return obs, info86 87 88def create_env(render_mode=None, block_xy=None, goal_xyz=None, environment = "FetchPickAndPlace-v3"):89    gym.register_envs(gymnasium_robotics)90 91    if(environment == "FetchReach-v3"):92        object = False93    else:94        object = True95 96    base_env = gym.make(environment, render_mode=render_mode)97    u = base_env.unwrapped98 99    # 1) compute table center in world coords100    #    – X,Y: same as the gripper’s initial XY (over table center)101    #    – Z: the table‐top height the wrapper uses (0.42 m)102    center_xy = u.initial_gripper_xpos[:2]        # e.g. [1.366, 0.750]103    table_z   = 0.42                              # match blk_qpos[2] in your wrapper104    table_center = np.array([*center_xy, table_z])105 106    # 2) turn your “relative” block_xy into an absolute XY107    if block_xy is not None:108        rel = np.array(block_xy, dtype=float)109        abs_block_xy = center_xy + rel110    else:111        abs_block_xy = None112 113    # 3) turn your “relative” goal_xyz into an absolute XYZ114    if goal_xyz is not None:115        rel = np.array(goal_xyz, dtype=float)116        abs_goal_xyz = table_center + rel117    else:118        abs_goal_xyz = None119 120    # 4) build the wrapped env with those absolutes121    env = CustomFetchWrapper(122        base_env,123        block_xy=abs_block_xy,124        goal_xyz=abs_goal_xyz,125        object=object126    )127    return env