Add files using upload-large-folder tool

Metaworld/metaworld/policies/policy.py

@@ -0,0 +1,71 @@
+import abc
+import warnings
+
+import numpy as np
+
+
+def assert_fully_parsed(func):
+    """Decorator function to ensure observations are fully parsed
+
+    Args:
+        func (Callable): The function to check
+
+    Returns:
+        (Callable): The input function, decorated to assert full parsing
+    """
+    def inner(obs):
+        obs_dict = func(obs)
+        assert len(obs) == sum(
+            [len(i) if isinstance(i, np.ndarray) else 1 for i in obs_dict.values()]
+        ), 'Observation not fully parsed'
+        return obs_dict
+    return inner
+
+
+def move(from_xyz, to_xyz, p):
+    """Computes action components that help move from 1 position to another
+
+    Args:
+        from_xyz (np.ndarray): The coordinates to move from (usually current position)
+        to_xyz (np.ndarray): The coordinates to move to
+        p (float): constant to scale response
+
+    Returns:
+        (np.ndarray): Response that will decrease abs(to_xyz - from_xyz)
+
+    """
+    error = to_xyz - from_xyz
+    response = p * error
+
+    if np.any(np.absolute(response) > 1.):
+        warnings.warn('Constant(s) may be too high. Environments clip response to [-1, 1]')
+
+    return response
+
+
+class Policy(abc.ABC):
+
+    @staticmethod
+    @abc.abstractmethod
+    def _parse_obs(obs):
+        """Pulls pertinent information out of observation and places in a dict.
+
+        Args:
+            obs (np.ndarray): Observation which conforms to env.observation_space
+
+        Returns:
+            dict: Dictionary which contains information from the observation
+        """
+        pass
+
+    @abc.abstractmethod
+    def get_action(self, obs):
+        """Gets an action in response to an observation.
+
+        Args:
+            obs (np.ndarray): Observation which conforms to env.observation_space
+
+        Returns:
+            np.ndarray: Array (usually 4 elements) representing the action to take
+        """
+        pass

Metaworld/metaworld/policies/sawyer_box_close_v1_policy.py

@@ -0,0 +1,60 @@
+import numpy as np
+
+from metaworld.policies.action import Action
+from metaworld.policies.policy import Policy, assert_fully_parsed, move
+
+
+class SawyerBoxCloseV1Policy(Policy):
+
+    @staticmethod
+    @assert_fully_parsed
+    def _parse_obs(obs):
+        return {
+            'hand_pos': obs[:3],
+            'lid_pos': obs[3:6],
+            'box_pos': obs[9:11],
+            'extra_info': obs[[6, 7, 8, 11]],
+        }
+
+    def get_action(self, obs):
+        o_d = self._parse_obs(obs)
+
+        action = Action({
+            'delta_pos': np.arange(3),
+            'grab_effort': 3
+        })
+
+        action['delta_pos'] = move(o_d['hand_pos'], to_xyz=self._desired_pos(o_d), p=25.)
+        action['grab_effort'] = self._grab_effort(o_d)
+
+        return action.array
+
+    @staticmethod
+    def _desired_pos(o_d):
+        pos_curr = o_d['hand_pos']
+        pos_lid = o_d['lid_pos'] + np.array([-.04, .0, -.06])
+        pos_box = np.array([*o_d['box_pos'], 0.15]) + np.array([-.04, .0, .0])
+
+        # If error in the XY plane is greater than 0.02, place end effector above the puck
+        if np.linalg.norm(pos_curr[:2] - pos_lid[:2]) > 0.01:
+            return pos_lid + np.array([0., 0., 0.1])
+        # Once XY error is low enough, drop end effector down on top of puck
+        elif abs(pos_curr[2] - pos_lid[2]) > 0.05:
+            return pos_lid
+        # If not at the same Z height as the goal, move up to that plane
+        elif abs(pos_curr[2] - pos_box[2]) > 0.04:
+            return np.array([pos_curr[0], pos_curr[1], pos_box[2]])
+        # Move to the goal
+        else:
+            return pos_box
+
+    @staticmethod
+    def _grab_effort(o_d):
+        pos_curr = o_d['hand_pos']
+        pos_puck = o_d['lid_pos'] + np.array([-.04, .0, -.06])
+
+        if np.linalg.norm(pos_curr[:2] - pos_puck[:2]) > 0.01 or abs(pos_curr[2] - pos_puck[2]) > 0.13:
+            return 0.
+        # While end effector is moving down toward the puck, begin closing the grabber
+        else:
+            return .8

Metaworld/metaworld/policies/sawyer_button_press_topdown_v2_policy.py

@@ -0,0 +1,40 @@
+import numpy as np
+
+from metaworld.policies.action import Action
+from metaworld.policies.policy import Policy, assert_fully_parsed, move
+
+
+class SawyerButtonPressTopdownV2Policy(Policy):
+
+    @staticmethod
+    @assert_fully_parsed
+    def _parse_obs(obs):
+        return {
+            'hand_pos': obs[:3],
+            'hand_closed': obs[3],
+            'button_pos': obs[4:7],
+            'unused_info': obs[7:],
+        }
+
+    def get_action(self, obs):
+        o_d = self._parse_obs(obs)
+
+        action = Action({
+            'delta_pos': np.arange(3),
+            'grab_effort': 3
+        })
+
+        action['delta_pos'] = move(o_d['hand_pos'], to_xyz=self._desired_pos(o_d), p=25.)
+        action['grab_effort'] = 1.
+
+        return action.array
+
+    @staticmethod
+    def _desired_pos(o_d):
+        pos_curr = o_d['hand_pos']
+        pos_button = o_d['button_pos']
+
+        if np.linalg.norm(pos_curr[:2] - pos_button[:2]) > 0.04:
+            return pos_button + np.array([0., 0., 0.1])
+        else:
+            return pos_button

Metaworld/metaworld/policies/sawyer_button_press_topdown_wall_v2_policy.py

@@ -0,0 +1,40 @@
+import numpy as np
+
+from metaworld.policies.action import Action
+from metaworld.policies.policy import Policy, assert_fully_parsed, move
+
+
+class SawyerButtonPressTopdownWallV2Policy(Policy):
+
+    @staticmethod
+    @assert_fully_parsed
+    def _parse_obs(obs):
+        return {
+            'hand_pos': obs[:3],
+            'hand_closed': obs[3],
+            'button_pos': obs[4:7],
+            'unused_info': obs[7:],
+        }
+
+    def get_action(self, obs):
+        o_d = self._parse_obs(obs)
+
+        action = Action({
+            'delta_pos': np.arange(3),
+            'grab_effort': 3
+        })
+
+        action['delta_pos'] = move(o_d['hand_pos'], to_xyz=self._desired_pos(o_d), p=25.)
+        action['grab_effort'] = -1.
+
+        return action.array
+
+    @staticmethod
+    def _desired_pos(o_d):
+        pos_curr = o_d['hand_pos']
+        pos_button = o_d['button_pos'] + np.array([.0, -.06, .0])
+
+        if np.linalg.norm(pos_curr[:2] - pos_button[:2]) > 0.04:
+            return pos_button + np.array([0., 0., 0.1])
+        else:
+            return pos_button

Metaworld/metaworld/policies/sawyer_button_press_v2_policy.py

@@ -0,0 +1,49 @@
+import numpy as np
+
+from metaworld.policies.action import Action
+from metaworld.policies.policy import Policy, move
+
+
+class SawyerButtonPressV2Policy(Policy):
+
+    @staticmethod
+    def _parse_obs(obs):
+        return {
+            'hand_pos': obs[:3],
+            'hand_closed': obs[3],
+            'button_pos': obs[4:7],
+            'unused_info': obs[7:],
+        }
+
+    def get_action(self, obs):
+        o_d = self._parse_obs(obs)
+
+        action = Action({
+            'delta_pos': np.arange(3),
+            'grab_effort': 3
+        })
+
+        action['delta_pos'] = move(o_d['hand_pos'], to_xyz=self.desired_pos(o_d), p=25.)
+        action['grab_effort'] = 0.
+
+        return action.array
+
+    @staticmethod
+    def desired_pos(o_d):
+        pos_curr = o_d['hand_pos']
+        pos_button = o_d['button_pos'] + np.array([0., 0., -0.07])
+        
+        # align the gripper with the button if the gripper does not have
+        # the same x and z position as the button.
+        hand_x, hand_y, hand_z = pos_curr
+        button_initial_x, button_initial_y, button_initial_z = pos_button
+        if not np.all(np.isclose(np.array([hand_x, hand_z]),
+                                np.array([button_initial_x, button_initial_z]),
+                                atol=0.02)):
+            pos_button[1] = pos_curr[1] - .1
+            return pos_button
+        # if the hand is aligned with the button, push the button in, by
+        # increasing the hand's y position
+        pos_button[1] += 0.02
+
+        return pos_button

Metaworld/metaworld/policies/sawyer_button_press_wall_v1_policy.py

@@ -0,0 +1,54 @@
+import numpy as np
+
+from metaworld.policies.action import Action
+from metaworld.policies.policy import Policy, move
+
+
+class SawyerButtonPressWallV1Policy(Policy):
+
+    @staticmethod
+    def _parse_obs(obs):
+        return {
+            'hand_pos': obs[:3],
+            'button_pos': obs[3:6],
+            'unused_info': obs[6:],
+        }
+
+    def get_action(self, obs):
+        o_d = self._parse_obs(obs)
+
+        action = Action({
+            'delta_pos': np.arange(3),
+            'grab_effort': 3
+        })
+
+        action['delta_pos'] = move(o_d['hand_pos'], to_xyz=self._desired_pos(o_d), p=15.)
+        action['grab_effort'] = self._grab_effort(o_d)
+
+        return action.array
+
+    @staticmethod
+    def _desired_pos(o_d):
+        pos_curr = o_d['hand_pos']
+        pos_button = o_d['button_pos'] + np.array([.0, .0, .04])
+
+        if abs(pos_curr[0] - pos_button[0]) > 0.02:
+            return np.array([pos_button[0], pos_curr[1], .3])
+        elif pos_button[1] - pos_curr[1] > 0.09:
+            return np.array([pos_button[0], pos_button[1], .3])
+        elif abs(pos_curr[2] - pos_button[2]) > 0.02:
+            return pos_button + np.array([.0, -.05, .0])
+        else:
+            return pos_button + np.array([.0, -.02, .0])
+
+    @staticmethod
+    def _grab_effort(o_d):
+        pos_curr = o_d['hand_pos']
+        pos_button = o_d['button_pos'] + np.array([.0, .0, .04])
+
+        if abs(pos_curr[0] - pos_button[0]) > 0.02 or \
+                pos_button[1] - pos_curr[1] > 0.09 or \
+                abs(pos_curr[2] - pos_button[2]) > 0.02:
+            return 1.
+        else:
+            return -1.

Metaworld/metaworld/policies/sawyer_button_press_wall_v2_policy.py

@@ -0,0 +1,55 @@
+import numpy as np
+
+from metaworld.policies.action import Action
+from metaworld.policies.policy import Policy, move
+
+
+class SawyerButtonPressWallV2Policy(Policy):
+
+    @staticmethod
+    def _parse_obs(obs):
+        return {
+            'hand_pos': obs[:3],
+            'hand_closed': obs[3],
+            'button_pos': obs[4:7],
+            'unused_info': obs[7:],
+        }
+
+    def get_action(self, obs):
+        o_d = self._parse_obs(obs)
+
+        action = Action({
+            'delta_pos': np.arange(3),
+            'grab_effort': 3
+        })
+
+        action['delta_pos'] = move(o_d['hand_pos'], to_xyz=self._desired_pos(o_d), p=15.)
+        action['grab_effort'] = self._grab_effort(o_d)
+
+        return action.array
+
+    @staticmethod
+    def _desired_pos(o_d):
+        pos_curr = o_d['hand_pos']
+        pos_button = o_d['button_pos'] + np.array([.0, .0, .04])
+
+        if abs(pos_curr[0] - pos_button[0]) > 0.02:
+            return np.array([pos_button[0], pos_curr[1], .3])
+        elif pos_button[1] - pos_curr[1] > 0.09:
+            return np.array([pos_button[0], pos_button[1], .3])
+        elif abs(pos_curr[2] - pos_button[2]) > 0.02:
+            return pos_button + np.array([.0, -.05, .0])
+        else:
+            return pos_button + np.array([.0, -.02, .0])
+
+    @staticmethod
+    def _grab_effort(o_d):
+        pos_curr = o_d['hand_pos']
+        pos_button = o_d['button_pos'] + np.array([.0, .0, .04])
+
+        if abs(pos_curr[0] - pos_button[0]) > 0.02 or \
+                pos_button[1] - pos_curr[1] > 0.09 or \
+                abs(pos_curr[2] - pos_button[2]) > 0.02:
+            return 1.
+        else:
+            return -1.

Metaworld/metaworld/policies/sawyer_coffee_button_v2_policy.py

@@ -0,0 +1,40 @@
+import numpy as np
+
+from metaworld.policies.action import Action
+from metaworld.policies.policy import Policy, assert_fully_parsed, move
+
+
+class SawyerCoffeeButtonV2Policy(Policy):
+
+    @staticmethod
+    @assert_fully_parsed
+    def _parse_obs(obs):
+        return {
+            'hand_pos': obs[:3],
+            'gripper': obs[3],
+            'button_pos': obs[4:7],
+            'unused_info': obs[7:],
+        }
+
+    def get_action(self, obs):
+        o_d = self._parse_obs(obs)
+
+        action = Action({
+            'delta_pos': np.arange(3),
+            'grab_effort': 3
+        })
+
+        action['delta_pos'] = move(o_d['hand_pos'], to_xyz=self._desired_pos(o_d), p=10.)
+        action['grab_effort'] = -1.
+
+        return action.array
+
+    @staticmethod
+    def _desired_pos(o_d):
+        pos_curr = o_d['hand_pos']
+        pos_button = o_d['button_pos'] + np.array([.0, .0, -.07])
+
+        if np.linalg.norm(pos_curr[[0, 2]] - pos_button[[0, 2]]) > 0.02:
+            return np.array([pos_button[0], pos_curr[1], pos_button[2]])
+        else:
+            return pos_button + np.array([.0, .2, .0])

Metaworld/metaworld/policies/sawyer_coffee_push_v2_policy.py

@@ -0,0 +1,56 @@
+import numpy as np
+
+from metaworld.policies.action import Action
+from metaworld.policies.policy import Policy, assert_fully_parsed, move
+
+
+class SawyerCoffeePushV2Policy(Policy):
+
+    @staticmethod
+    @assert_fully_parsed
+    def _parse_obs(obs):
+        return {
+            'hand_pos': obs[:3],
+            'gripper': obs[3],
+            'mug_pos': obs[4:7],
+            'goal_xy': obs[-3:-1],
+            'unused_info_1': obs[7:-3],
+            'unused_info_2': obs[-1],
+        }
+
+    def get_action(self, obs):
+        o_d = self._parse_obs(obs)
+
+        action = Action({
+            'delta_pos': np.arange(3),
+            'grab_effort': 3
+        })
+
+        action['delta_pos'] = move(o_d['hand_pos'], to_xyz=self._desired_pos(o_d), p=10.)
+        action['grab_effort'] = self._grab_effort(o_d)
+
+        return action.array
+
+    @staticmethod
+    def _desired_pos(o_d):
+        pos_curr = o_d['hand_pos']
+        pos_mug = o_d['mug_pos'] + np.array([.01, .0, .05])
+        pos_goal = o_d['goal_xy']
+
+        if np.linalg.norm(pos_curr[:2] - pos_mug[:2]) > 0.06:
+            return pos_mug + np.array([.0, .0, .2])
+        elif abs(pos_curr[2] - pos_mug[2]) > 0.02:
+            return pos_mug
+        else:
+            return np.array([pos_goal[0], pos_goal[1], .1])
+
+    @staticmethod
+    def _grab_effort(o_d):
+        pos_curr = o_d['hand_pos']
+        pos_mug = o_d['mug_pos'] + np.array([.01, .0, .05])
+
+        if np.linalg.norm(pos_curr[:2] - pos_mug[:2]) > 0.06 or \
+                abs(pos_curr[2] - pos_mug[2]) > 0.1:
+            return -1.
+        else:
+            return .5

Metaworld/metaworld/policies/sawyer_dial_turn_v1_policy.py

@@ -0,0 +1,39 @@
+import numpy as np
+
+from metaworld.policies.action import Action
+from metaworld.policies.policy import Policy, assert_fully_parsed, move
+
+
+class SawyerDialTurnV1Policy(Policy):
+
+    @staticmethod
+    @assert_fully_parsed
+    def _parse_obs(obs):
+        return {
+            'hand_pos': obs[:3],
+            'dial_pos': obs[3:6],
+            'goal_pos': obs[6:],
+        }
+
+    def get_action(self, obs):
+        o_d = self._parse_obs(obs)
+
+        action = Action({
+            'delta_pos': np.arange(3),
+            'grab_pow': 3
+        })
+
+        action['delta_pos'] = move(o_d['hand_pos'], to_xyz=self._desired_xyz(o_d), p=5.)
+        action['grab_pow'] = 0.
+
+        return action.array
+
+    @staticmethod
+    def _desired_xyz(o_d):
+        hand_pos = o_d['hand_pos']
+        dial_pos = o_d['dial_pos'] + np.array([0.0, -0.028, 0.0])
+        if abs(hand_pos[2] - dial_pos[2]) > 0.02:
+            return np.array([hand_pos[0], hand_pos[1], dial_pos[2]])
+        elif abs(hand_pos[1] - dial_pos[1]) > 0.02:
+            return np.array([dial_pos[0]+0.20, dial_pos[1], dial_pos[2]])
+        return np.array([dial_pos[0]-0.10, dial_pos[1], dial_pos[2]])

Metaworld/metaworld/policies/sawyer_dial_turn_v2_policy.py

@@ -0,0 +1,41 @@
+import numpy as np
+
+from metaworld.policies.action import Action
+from metaworld.policies.policy import Policy, assert_fully_parsed, move
+
+
+class SawyerDialTurnV2Policy(Policy):
+
+    @staticmethod
+    @assert_fully_parsed
+    def _parse_obs(obs):
+        return {
+            'hand_pos': obs[:3],
+            'unused_gripper_open': obs[3],
+            'dial_pos': obs[4:7],
+            'extra_info': obs[7:],
+        }
+
+    def get_action(self, obs):
+        o_d = self._parse_obs(obs)
+
+        action = Action({
+            'delta_pos': np.arange(3),
+            'grab_pow': 3
+        })
+
+        action['delta_pos'] = move(o_d['hand_pos'], to_xyz=self._desired_pos(o_d), p=10.)
+        action['grab_pow'] = 1.
+
+        return action.array
+
+    @staticmethod
+    def _desired_pos(o_d):
+        hand_pos = o_d['hand_pos']
+        dial_pos = o_d['dial_pos'] + np.array([0.05, 0.02, 0.09])
+
+        if np.linalg.norm(hand_pos[:2] - dial_pos[:2]) > 0.02:
+            return np.array([*dial_pos[:2], 0.2])
+        if abs(hand_pos[2] - dial_pos[2]) > 0.02:
+            return dial_pos
+        return dial_pos + np.array([-.05, .005, .0])

Metaworld/metaworld/policies/sawyer_disassemble_v1_policy.py

@@ -0,0 +1,59 @@
+import numpy as np
+
+from metaworld.policies.action import Action
+from metaworld.policies.policy import Policy, assert_fully_parsed, move
+
+
+class SawyerDisassembleV1Policy(Policy):
+
+    @staticmethod
+    @assert_fully_parsed
+    def _parse_obs(obs):
+        return {
+            'hand_pos': obs[:3],
+            'wrench_pos': obs[3:6],
+            'peg_pos': obs[9:],
+            'unused_info': obs[6:9],
+        }
+
+    def get_action(self, obs):
+        o_d = self._parse_obs(obs)
+
+        action = Action({
+            'delta_pos': np.arange(3),
+            'grab_effort': 3
+        })
+
+        action['delta_pos'] = move(o_d['hand_pos'], to_xyz=self._desired_pos(o_d), p=10.)
+        action['grab_effort'] = self._grab_effort(o_d)
+
+        return action.array
+
+    @staticmethod
+    def _desired_pos(o_d):
+        pos_curr = o_d['hand_pos']
+        pos_wrench = o_d['wrench_pos'] + np.array([.01, -.01, .01])
+        pos_peg = o_d['peg_pos'] + np.array([.07, .0, .15])
+
+        # If XY error is greater than 0.02, place end effector above the wrench
+        if np.linalg.norm(pos_curr[:2] - pos_wrench[:2]) > 0.02:
+            return pos_wrench + np.array([0., 0., 0.12])
+        # Once XY error is low enough, drop end effector down on top of wrench
+        elif abs(pos_curr[2] - pos_wrench[2]) > 0.03:
+            return pos_wrench
+        # If still hooked on peg, move upwards
+        elif pos_wrench[2] < 0.12:
+            return pos_peg + np.array([.0, .0, .1])
+        # Move away from peg
+        else:
+            return pos_curr + np.array([.0, -.1, .0])
+
+    @staticmethod
+    def _grab_effort(o_d):
+        pos_curr = o_d['hand_pos']
+        pos_wrench = o_d['wrench_pos'] + np.array([.01, .0, .0])
+
+        if np.linalg.norm(pos_curr[:2] - pos_wrench[:2]) > 0.02 or abs(pos_curr[2] - pos_wrench[2]) > 0.08:
+            return 0.
+        else:
+            return 0.8

Metaworld/metaworld/policies/sawyer_disassemble_v2_policy.py

@@ -0,0 +1,57 @@
+import numpy as np
+
+from metaworld.policies.action import Action
+from metaworld.policies.policy import Policy, assert_fully_parsed, move
+
+
+class SawyerDisassembleV2Policy(Policy):
+
+    @staticmethod
+    @assert_fully_parsed
+    def _parse_obs(obs):
+        return {
+            'hand_pos': obs[:3],
+            'gripper': obs[3],
+            'wrench_pos': obs[4:7],
+            'peg_pos': obs[-3:],
+            'unused_info': obs[7:-3],
+        }
+
+    def get_action(self, obs):
+        o_d = self._parse_obs(obs)
+
+        action = Action({
+            'delta_pos': np.arange(3),
+            'grab_effort': 3
+        })
+
+        action['delta_pos'] = move(o_d['hand_pos'], to_xyz=self._desired_pos(o_d), p=10.)
+        action['grab_effort'] = self._grab_effort(o_d)
+
+        return action.array
+
+    @staticmethod
+    def _desired_pos(o_d):
+        pos_curr = o_d['hand_pos']
+        pos_wrench = o_d['wrench_pos'] + np.array([-.02, .0, .01])
+        pos_peg = o_d['peg_pos'] + np.array([.12, .0, .14])
+
+        # If XY error is greater than 0.02, place end effector above the wrench
+        if np.linalg.norm(pos_curr[:2] - pos_wrench[:2]) > 0.02:
+            return pos_wrench + np.array([0., 0., 0.1])
+        # Once XY error is low enough, drop end effector down on top of wrench
+        elif abs(pos_curr[2] - pos_wrench[2]) > 0.03:
+            return pos_wrench
+        # Move upwards
+        else:
+            return pos_curr + np.array([.0, .0, .1])
+
+    @staticmethod
+    def _grab_effort(o_d):
+        pos_curr = o_d['hand_pos']
+        pos_wrench = o_d['wrench_pos'] + np.array([-.02, .0, .01])
+
+        if np.linalg.norm(pos_curr[:2] - pos_wrench[:2]) > 0.02 or abs(pos_curr[2] - pos_wrench[2]) > 0.07:
+            return 0.
+        else:
+            return 0.8

Metaworld/metaworld/policies/sawyer_door_close_v2_policy.py

@@ -0,0 +1,54 @@
+import numpy as np
+
+from metaworld.policies.action import Action
+from metaworld.policies.policy import Policy, assert_fully_parsed, move
+
+
+class SawyerDoorCloseV2Policy(Policy):
+
+    @staticmethod
+    @assert_fully_parsed
+    def _parse_obs(obs):
+        return {
+            'hand_pos': obs[:3],
+            'unused_1': obs[3],
+            'door_pos': obs[4:7],
+            'unused_2': obs[7:-3],
+            'goal_pos': obs[-3:],
+        }
+
+    def get_action(self, obs):
+        o_d = self._parse_obs(obs)
+
+        action = Action({
+            'delta_pos': np.arange(3),
+            'grab_effort': 3
+        })
+
+        action['delta_pos'] = move(o_d['hand_pos'], to_xyz=self._desired_pos(o_d), p=25.)
+        action['grab_effort'] = 1.
+
+        return action.array
+
+    @staticmethod
+    def _desired_pos(o_d):
+        pos_curr = o_d['hand_pos']
+        pos_door = o_d['door_pos']
+        pos_door += np.array([0.05, 0.12, 0.1])
+        pos_goal = o_d['goal_pos']
+
+        # # if to the right of door handle///
+        # if pos_curr[0] > pos_door[0]:
+        #     # if below door handle by more than 0.2
+        #     if pos_curr[2] < pos_door[2] + 0.2:
+        #         # rise above door handle by ~0.2
+        #         return np.array([pos_curr[0], pos_curr[1], pos_door[2] + 0.25])
+        #     else:
+        #         # move toward door handle in XY plane
+        #         return np.array([pos_door[0] - 0.02, pos_door[1], pos_curr[2]])
+        # # put end effector on the outer edge of door handle (still above it)
+        # elif abs(pos_curr[2] - pos_door[2]) > 0.04:
+        #     return pos_door + np.array([-0.02, 0., 0.])
+        # # push from outer edge toward door handle's centroid
+        # else:
+        return pos_goal

Metaworld/metaworld/policies/sawyer_drawer_close_v1_policy.py

@@ -0,0 +1,49 @@
+import numpy as np
+
+from metaworld.policies.action import Action
+from metaworld.policies.policy import Policy, assert_fully_parsed, move
+
+
+class SawyerDrawerCloseV1Policy(Policy):
+
+    @staticmethod
+    @assert_fully_parsed
+    def _parse_obs(obs):
+        return {
+            'hand_pos': obs[:3],
+            'drwr_pos': obs[3:6],
+            'unused_info': obs[6:],
+        }
+
+    def get_action(self, obs):
+        o_d = self._parse_obs(obs)
+
+        action = Action({
+            'delta_pos': np.arange(3),
+            'grab_effort': 3
+        })
+
+        action['delta_pos'] = move(o_d['hand_pos'], to_xyz=self._desired_pos(o_d), p=10.)
+        action['grab_effort'] = 1.
+
+        return action.array
+
+    @staticmethod
+    def _desired_pos(o_d):
+        pos_curr = o_d['hand_pos']
+        pos_drwr = o_d['drwr_pos']
+
+        # if further forward than the drawer...
+        if pos_curr[1] > pos_drwr[1]:
+            if pos_curr[2] < pos_drwr[2] + 0.4:
+                # rise up quickly (Z direction)
+                return np.array([pos_curr[0], pos_curr[1], pos_drwr[2] + 0.5])
+            else:
+                # move to front edge of drawer handle, but stay high in Z
+                return pos_drwr + np.array([0., -0.075, 0.4])
+        # drop down to touch drawer handle
+        elif abs(pos_curr[2] - pos_drwr[2]) > 0.04:
+            return pos_drwr + np.array([0., -0.075, 0.])
+        # push toward drawer handle's centroid
+        else:
+            return pos_drwr

Metaworld/metaworld/policies/sawyer_drawer_close_v2_policy.py

@@ -0,0 +1,50 @@
+import numpy as np
+
+from metaworld.policies.action import Action
+from metaworld.policies.policy import Policy, assert_fully_parsed, move
+
+
+class SawyerDrawerCloseV2Policy(Policy):
+
+    @staticmethod
+    @assert_fully_parsed
+    def _parse_obs(obs):
+        return {
+            'hand_pos': obs[:3],
+            'unused_grasp_info': obs[3],
+            'drwr_pos': obs[4:7],
+            'unused_info': obs[7:],
+        }
+
+    def get_action(self, obs):
+        o_d = self._parse_obs(obs)
+
+        action = Action({
+            'delta_pos': np.arange(3),
+            'grab_effort': 3
+        })
+
+        action['delta_pos'] = move(o_d['hand_pos'], to_xyz=self._desired_pos(o_d), p=25.)
+        action['grab_effort'] = 1.
+
+        return action.array
+
+    @staticmethod
+    def _desired_pos(o_d):
+        pos_curr = o_d['hand_pos']
+        pos_drwr = o_d['drwr_pos'] + np.array([.0, .0, -.02])
+
+        # if further forward than the drawer...
+        if pos_curr[1] > pos_drwr[1]:
+            if pos_curr[2] < pos_drwr[2] + 0.23:
+                # rise up quickly (Z direction)
+                return np.array([pos_curr[0], pos_curr[1], pos_drwr[2] + 0.5])
+            else:
+                # move to front edge of drawer handle, but stay high in Z
+                return pos_drwr + np.array([0., -0.075, 0.23])
+        # drop down to touch drawer handle
+        elif abs(pos_curr[2] - pos_drwr[2]) > 0.04:
+            return pos_drwr + np.array([0., -0.075, 0.])
+        # push toward drawer handle's centroid
+        else:
+            return pos_drwr

Metaworld/metaworld/policies/sawyer_drawer_open_v1_policy.py

@@ -0,0 +1,48 @@
+import numpy as np
+
+from metaworld.policies.action import Action
+from metaworld.policies.policy import Policy, assert_fully_parsed, move
+
+
+class SawyerDrawerOpenV1Policy(Policy):
+
+    @staticmethod
+    @assert_fully_parsed
+    def _parse_obs(obs):
+        return {
+            'hand_pos': obs[:3],
+            'drwr_pos': obs[3:6],
+            'unused_info': obs[6:],
+        }
+
+    def get_action(self, obs):
+        o_d = self._parse_obs(obs)
+
+        action = Action({
+            'delta_pos': np.arange(3),
+            'grab_effort': 3
+        })
+
+        # NOTE this policy looks different from the others because it must
+        # modify its p constant part-way through the task
+        pos_curr = o_d['hand_pos']
+        pos_drwr = o_d['drwr_pos']
+
+        # align end effector's Z axis with drawer handle's Z axis
+        if np.linalg.norm(pos_curr[:2] - pos_drwr[:2]) > 0.06:
+            to_pos = pos_drwr + np.array([0., 0., 0.3])
+            action['delta_pos'] = move(o_d['hand_pos'], to_pos, p=4.)
+        # drop down to touch drawer handle
+        elif abs(pos_curr[2] - pos_drwr[2]) > 0.04:
+            to_pos = pos_drwr
+            action['delta_pos'] = move(o_d['hand_pos'], to_pos, p=4.)
+        # push toward a point just behind the drawer handle
+        # also increase p value to apply more force
+        else:
+            to_pos = pos_drwr + np.array([0., -0.06, 0.])
+            action['delta_pos'] = move(o_d['hand_pos'], to_pos, p=50.)
+
+        # keep gripper open
+        action['grab_effort'] = -1.
+
+        return action.array

Metaworld/metaworld/policies/sawyer_faucet_close_v1_policy.py

@@ -0,0 +1,41 @@
+import numpy as np
+
+from metaworld.policies.action import Action
+from metaworld.policies.policy import Policy, assert_fully_parsed, move
+
+
+class SawyerFaucetCloseV1Policy(Policy):
+    
+    @staticmethod
+    @assert_fully_parsed
+    def _parse_obs(obs):
+        return {
+            'hand_pos': obs[:3],
+            'faucet_pos': obs[3:6],
+            'unused_info': obs[6:],
+        }
+
+    def get_action(self, obs):
+        o_d = self._parse_obs(obs)
+
+        action = Action({
+            'delta_pos': np.arange(3),
+            'grab_effort': 3
+        })
+
+        action['delta_pos'] = move(o_d['hand_pos'], to_xyz=self._desired_pos(o_d), p=25.)
+        action['grab_effort'] = 1.
+
+        return action.array
+
+    @staticmethod
+    def _desired_pos(o_d):
+        pos_curr = o_d['hand_pos']
+        pos_faucet = o_d['faucet_pos'] + np.array([.02, .0, .0])
+
+        if np.linalg.norm(pos_curr[:2] - pos_faucet[:2]) > 0.04:
+            return pos_faucet + np.array([.0, .0, .1])
+        elif abs(pos_curr[2] - pos_faucet[2]) > 0.04:
+            return pos_faucet
+        else:
+            return pos_faucet + np.array([-.1, .05, .0])

Metaworld/metaworld/policies/sawyer_faucet_open_v2_policy.py

@@ -0,0 +1,42 @@
+import numpy as np
+
+from metaworld.policies.action import Action
+from metaworld.policies.policy import Policy, assert_fully_parsed, move
+
+
+class SawyerFaucetOpenV2Policy(Policy):
+
+    @staticmethod
+    @assert_fully_parsed
+    def _parse_obs(obs):
+        return {
+            'hand_pos': obs[:3],
+            'unused_gripper': obs[3],
+            'faucet_pos': obs[4:7],
+            'unused_info': obs[7:],
+        }
+
+    def get_action(self, obs):
+        o_d = self._parse_obs(obs)
+
+        action = Action({
+            'delta_pos': np.arange(3),
+            'grab_effort': 3
+        })
+
+        action['delta_pos'] = move(o_d['hand_pos'], to_xyz=self._desired_pos(o_d), p=25.)
+        action['grab_effort'] = 1.
+
+        return action.array
+
+    @staticmethod
+    def _desired_pos(o_d):
+        pos_curr = o_d['hand_pos']
+        pos_faucet = o_d['faucet_pos'] + np.array([-.04, .0, .03])
+
+        if np.linalg.norm(pos_curr[:2] - pos_faucet[:2]) > 0.04:
+            return pos_faucet + np.array([.0, .0, .1])
+        elif abs(pos_curr[2] - pos_faucet[2]) > 0.04:
+            return pos_faucet
+        else:
+            return pos_faucet + np.array([.1, .05, .0])

Metaworld/metaworld/policies/sawyer_hand_insert_v2_policy.py

@@ -0,0 +1,59 @@
+import numpy as np
+
+from metaworld.policies.action import Action
+from metaworld.policies.policy import Policy, assert_fully_parsed, move
+
+
+class SawyerHandInsertV2Policy(Policy):
+
+    @staticmethod
+    @assert_fully_parsed
+    def _parse_obs(obs):
+        return {
+            'hand_pos': obs[:3],
+            'gripper': obs[3],
+            'obj_pos': obs[4:7],
+            'goal_pos': obs[-3:],
+            'unused_info': obs[7:-3],
+        }
+
+    def get_action(self, obs):
+        o_d = self._parse_obs(obs)
+
+        action = Action({
+            'delta_pos': np.arange(3),
+            'grab_effort': 3
+        })
+
+        action['delta_pos'] = move(o_d['hand_pos'], to_xyz=self._desired_pos(o_d), p=10.)
+        action['grab_effort'] = self._grab_effort(o_d)
+
+        return action.array
+
+    @staticmethod
+    def _desired_pos(o_d):
+        hand_pos = o_d['hand_pos']
+        obj_pos = o_d['obj_pos']
+        goal_pos = o_d['goal_pos']
+
+        # If error in the XY plane is greater than 0.02, place end effector above the puck
+        if np.linalg.norm(hand_pos[:2] - obj_pos[:2]) > 0.02:
+            return obj_pos + np.array([0., 0., 0.1])
+        # Once XY error is low enough, drop end effector down on top of puck
+        elif abs(hand_pos[2] - obj_pos[2]) > 0.05:
+            return obj_pos + np.array([0., 0., 0.03])
+        # If not above goal, move to be directly above goal
+        elif np.linalg.norm(hand_pos[:2] - goal_pos[:2]) > 0.04:
+            return np.array([goal_pos[0], goal_pos[1], hand_pos[2]])
+        else:
+            return goal_pos
+
+    @staticmethod
+    def _grab_effort(o_d):
+        hand_pos = o_d['hand_pos']
+        obj_pos = o_d['obj_pos']
+
+        if np.linalg.norm(hand_pos[:2] - obj_pos[:2]) > 0.02 or abs(hand_pos[2] - obj_pos[2]) > 0.1:
+            return 0.
+        else:
+            return 0.65

Metaworld/metaworld/policies/sawyer_handle_press_side_v2_policy.py

@@ -0,0 +1,40 @@
+import numpy as np
+
+from metaworld.policies.action import Action
+from metaworld.policies.policy import Policy, assert_fully_parsed, move
+
+
+class SawyerHandlePressSideV2Policy(Policy):
+
+    @staticmethod
+    @assert_fully_parsed
+    def _parse_obs(obs):
+        return {
+            'hand_pos': obs[:3],
+            'gripper': obs[3],
+            'handle_pos': obs[4:7],
+            'unused_info': obs[7:],
+        }
+
+    def get_action(self, obs):
+        o_d = self._parse_obs(obs)
+
+        action = Action({
+            'delta_pos': np.arange(3),
+            'grab_effort': 3
+        })
+
+        action['delta_pos'] = move(o_d['hand_pos'], to_xyz=self._desired_pos(o_d), p=25.)
+        action['grab_effort'] = 1.
+
+        return action.array
+
+    @staticmethod
+    def _desired_pos(o_d):
+        pos_curr = o_d['hand_pos']
+        pos_button = o_d['handle_pos']
+
+        if np.linalg.norm(pos_curr[:2] - pos_button[:2]) > 0.02:
+            return pos_button + np.array([0., 0., 0.2])
+        else:
+            return pos_button + np.array([.0, .0, -.5])

Metaworld/metaworld/policies/sawyer_handle_pull_side_v1_policy.py

@@ -0,0 +1,43 @@
+import numpy as np
+
+from metaworld.policies.action import Action
+from metaworld.policies.policy import Policy, assert_fully_parsed, move
+
+
+class SawyerHandlePullSideV1Policy(Policy):
+
+    @staticmethod
+    @assert_fully_parsed
+    def _parse_obs(obs):
+        return {
+            'hand_pos': obs[:3],
+            'handle_pos': obs[3:6],
+            'unused_info': obs[6:],
+        }
+
+    def get_action(self, obs):
+        o_d = self._parse_obs(obs)
+
+        action = Action({
+            'delta_pos': np.arange(3),
+            'grab_effort': 3
+        })
+
+        action['delta_pos'] = move(o_d['hand_pos'], to_xyz=self._desired_pos(o_d), p=25.)
+        action['grab_effort'] = 1.
+
+        return action.array
+
+    @staticmethod
+    def _desired_pos(o_d):
+        pos_curr = o_d['hand_pos']
+        pos_button = o_d['handle_pos'] + np.array([.02, .0, .0])
+
+        if abs(pos_curr[1] - pos_button[1]) > 0.04:
+            return pos_button + np.array([0., 0., 0.2])
+        elif abs(pos_curr[2] - pos_button[2]) > 0.03:
+            return pos_button + np.array([.1, .0, -.01])
+        elif abs(pos_curr[0] - pos_button[0]) > .01:
+            return np.array([pos_button[0] - .04, pos_button[1], pos_curr[2]])
+        else:
+            return pos_button + np.array([-.04, .0, .1])