<commit_before>"""This environment allows a Sawyer to interact with multiple custom objects.
It is largely copied over from <NAME> and <NAME>:
https://github.com/febert/visual_mpc/ but adapted to the multiworld interface.
"""

import multiworld
import numpy as np
import cv2
import mujoco_py
from pyquaternion import Quaternion
from multiworld.envs.mujoco.util.create_xml import create_object_xml, create_root_xml, clean_xml
from collections import OrderedDict
from multiworld.envs.mujoco.util.interpolation import TwoPointCSpline

from mujoco_py.builder import MujocoException
import copy
from multiworld.envs.env_util import get_stat_in_paths, \
    create_stats_ordered_dict, get_asset_full_path
from multiworld.core.multitask_env import MultitaskEnv
from multiworld.envs.mujoco.sawyer_xyz.base import SawyerXYZEnv

from gym.spaces import Box, Dict
import time

from multiworld.core.image_env import ImageEnv
from multiworld.envs.mujoco.cameras import sawyer_pusher_camera_upright_v2

def quat_to_zangle(quat):
    angle = -(Quaternion(axis = [0,1,0], angle = np.pi).inverse * Quaternion(quat)).angle
    if angle < 0:
        return angle + 2 * np.pi
    return angle

def zangle_to_quat(zangle):
    """
    :param zangle in rad
    :return: quaternion
    """
    return (Quaternion(axis=[0,1,0], angle=np.pi) * Quaternion(axis=[0, 0, -1], angle= zangle)).elements

BASE_DIR = '/'.join(str.split(multiworld.__file__, '/')[:-2])
asset_base_path = BASE_DIR + '/multiworld/envs/assets/multi_object_sawyer_xyz/'

low_bound = np.array([-0.27, 0.52, 0.05, 0, -1])
high_bound = np.array([0.27, 0.95, 0.1, 2 * np.pi - 0.001, 1])
NEUTRAL_JOINTS = np.array([1.65474475, - 0.53312487, - 0.65980174, 1.1841825, 0.62772584, 1.11682223, 1.31015104, -0.05, 0.05])

class MultiSawyerEnv(MultitaskEnv, SawyerXYZEnv):
    def __init__(self,
        filename='sawyer_multiobj.xml',
        mode_rel=np.array([True, True, True, True, False]),
        num_objects = 3, object_mass = 1, friction=1.0, finger_sensors=True,
        maxlen=0.12, minlen=0.01, preload_obj_dict=None,
        object_meshes=['Bowl', 'GlassBowl', 'LotusBowl01', 'ElephantBowl', 'RuggedBowl'],
        obj_classname = 'freejoint', block_height=0.02, block_width = 0.02,
        viewer_image_height = 84, viewer_image_width = 84, skip_first=100,
        substeps=100, init_hand_xyz=(0, 0.4, 0.07),
        randomize_initial_pos = False, state_goal = None,
        randomize_goal_at_reset = False, fix_rotation = True,
        fix_reset_pos = True, fix_gripper = True, do_render = False,
        match_orientation = False,
    ):
        self.quick_init(locals())

        base_filename = asset_base_path + filename

        friction_params = (friction, 0.1, 0.02)
        self.obj_stat_prop = create_object_xml(base_filename, num_objects, object_mass,
                                               friction_params, object_meshes, finger_sensors,
                                               maxlen, minlen, preload_obj_dict, obj_classname,
                                               block_height, block_width)


        gen_xml = create_root_xml(base_filename)
        SawyerXYZEnv.__init__(self, model_name=gen_xml)
        clean_xml(gen_xml)

        if self.sim.model.nmocap > 0 and self.sim.model.eq_data is not None:
            for i in range(self.sim.model.eq_data.shape[0]):
                if self.sim.model.eq_type[i] == mujoco_py.const.EQ_WELD:
                    # Define the xyz + quat of the mocap relative to the hand
                    self.sim.model.eq_data[i, :] = np.array(
                        [0., 0., 0., 1., 0., 0., 0.]
                    )

        self._base_sdim, self._base_adim, self.mode_rel = 5, 5, mode_rel
        self.hand_low = np.array((-0.03, 0.55, 0.05))
        self.hand_high = np.array((-0.0, 0.65, 0.2))
        self.mocap_high = self.hand_high
        self.mocap_low = self.hand_low
        self.action_scale = 1.0/100
        self.num_objects, self.skip_first, self.substeps = num_objects, skip_first, substeps
        self.randomize_initial_pos = randomize_initial_pos
        self.finger_sensors, self._maxlen = finger_sensors, maxlen
        self._threshold = 0
        self._previous_target_qpos, self._n_joints = None, 8
        self._object_pos = np.zeros((num_objects, 7))
        self._reset_xyz = np.zeros((num_objects, 3))
        self._reset_quat = np.zeros((num_objects, 4))
        self._reset_pos = np.copy(self._object_pos)
        self._object_names = ['obj_' + str(i) for i in range(num_objects)]
        self._initialized = False
        self._state_goal = state_goal
        self.fix_rotation = fix_rotation
        self.fix_gripper = fix_gripper
        self._randomize_goal_at_reset = randomize_goal_at_reset
        self.do_render = do_render
        self.init_hand_xyz = np.array(init_hand_xyz)
        self.match_orientation = match_orientation
        self.goal_dim_per_object = 7 if match_orientation else 3

        # self.finger_sensors = False # turning this off for now
        self.action_space = Box(np.array([-1, -1, -1, -1, -1]),
                        np.array([1, 1, 1, 1, 1]), dtype=np.float32)
        # action space is 5-dim: x, y, z, ? ?
        self.hand_space = Box(self.hand_low, self.hand_high, dtype=np.float32)
        obs_dim = 7 * num_objects + 3
        obs_space = Box(np.array([-1] * obs_dim), np.array([1] * obs_dim), dtype=np.float32)
        goal_space = Box(
            np.array([-1] * self.goal_dim_per_object * num_objects),
            np.array([1] * self.goal_dim_per_object * num_objects),
            dtype=np.float32,
        )
        self.observation_space = Dict([
            ('observation', obs_space),
            ('desired_goal', goal_space),
            ('achieved_goal', goal_space),
            ('state_observation', obs_space),
            ('state_desired_goal', goal_space),
            ('state_achieved_goal', goal_space),
        ])

        self._state_goal = self.sample_goal()
        self.fix_reset_pos = fix_reset_pos
        o = self.reset()

    def _clip_gripper(self):
        self.sim.data.qpos[7:9] = np.clip(self.sim.data.qpos[7:9], [-0.055, 0.0027], [-0.0027, 0.055])

    def samp_xyz_rot(self):
        rand_xyz = np.random.uniform(low_bound[:3] + self._maxlen / 2 + 0.02, high_bound[:3] - self._maxlen / 2 + 0.02)
        rand_xyz[-1] = 0.05
        return rand_xyz, np.random.uniform(-np.pi / 2, np.pi / 2)

    def project_point(self, point, camera):
        model_matrix = np.zeros((4, 4))
        model_matrix[:3, :3] = self.sim.data.get_camera_xmat(camera).T
        model_matrix[-1, -1] = 1

        fovy_radians = np.deg2rad(self.sim.model.cam_fovy[self.sim.model.camera_name2id(camera)])
        uh = 1. / np.tan(fovy_radians / 2)
        uw = uh / (self._frame_width / self._frame_height)
        extent = self.sim.model.stat.extent
        far, near = self.sim.model.vis.map.zfar * extent, self.sim.model.vis.map.znear * extent
        view_matrix = np.array([[uw, 0., 0., 0.],                        #matrix definition from
                                [0., uh, 0., 0.],                        #https://stackoverflow.com/questions/18404890/how-to-build-perspective-projection-matrix-no-api
                                [0., 0., far / (far - near), -1.],
                                [0., 0., -2*far*near/(far - near), 0.]]) #Note Mujoco doubles this quantity

        MVP_matrix = view_matrix.dot(model_matrix)
        world_coord = np.ones((4, 1))
        world_coord[:3, 0] = point - self.sim.data.get_camera_xpos(camera)

        clip = MVP_matrix.dot(world_coord)
        ndc = clip[:3] / clip[3]  # everything should now be in -1 to 1!!
        col, row = (ndc[0] + 1) * self._frame_width / 2, (-ndc[1] + 1) * self._frame_height / 2

        return self._frame_height - row, col                 #rendering flipped around in height

    def get_desig_pix(self, cams, target_width, round=True):
        qpos_dim = self._n_joints      # the states contains pos and vel
        assert self.sim.data.qpos.shape[0] == qpos_dim + 7 * self.num_objects<commit_msg>Add bounds as params<commit_after>"""This environment allows a Sawyer to interact with multiple custom objects.
It is largely copied over from <NAME> and <NAME>:
https://github.com/febert/visual_mpc/ but adapted to the multiworld interface.
"""

import multiworld
import numpy as np
import cv2
import mujoco_py
from pyquaternion import Quaternion
from multiworld.envs.mujoco.util.create_xml import create_object_xml, create_root_xml, clean_xml
from collections import OrderedDict
from multiworld.envs.mujoco.util.interpolation import TwoPointCSpline

from mujoco_py.builder import MujocoException
import copy
from multiworld.envs.env_util import get_stat_in_paths, \
    create_stats_ordered_dict, get_asset_full_path
from multiworld.core.multitask_env import MultitaskEnv
from multiworld.envs.mujoco.sawyer_xyz.base import SawyerXYZEnv

from gym.spaces import Box, Dict
import time

from multiworld.core.image_env import ImageEnv
from multiworld.envs.mujoco.cameras import sawyer_pusher_camera_upright_v2

def quat_to_zangle(quat):
    angle = -(Quaternion(axis = [0,1,0], angle = np.pi).inverse * Quaternion(quat)).angle
    if angle < 0:
        return angle + 2 * np.pi
    return angle

def zangle_to_quat(zangle):
    """
    :param zangle in rad
    :return: quaternion
    """
    return (Quaternion(axis=[0,1,0], angle=np.pi) * Quaternion(axis=[0, 0, -1], angle= zangle)).elements

BASE_DIR = '/'.join(str.split(multiworld.__file__, '/')[:-2])
asset_base_path = BASE_DIR + '/multiworld/envs/assets/multi_object_sawyer_xyz/'

low_bound = np.array([-0.27, 0.52, 0.05, 0, -1])
high_bound = np.array([0.27, 0.95, 0.1, 2 * np.pi - 0.001, 1])
NEUTRAL_JOINTS = np.array([1.65474475, - 0.53312487, - 0.65980174, 1.1841825, 0.62772584, 1.11682223, 1.31015104, -0.05, 0.05])

class MultiSawyerEnv(MultitaskEnv, SawyerXYZEnv):
    def __init__(self,
        filename='sawyer_multiobj.xml',
        mode_rel=np.array([True, True, True, True, False]),
        num_objects = 3, object_mass = 1, friction=1.0, finger_sensors=True,
        maxlen=0.12, minlen=0.01, preload_obj_dict=None,
        object_meshes=['Bowl', 'GlassBowl', 'LotusBowl01', 'ElephantBowl', 'RuggedBowl'],
        obj_classname = 'freejoint', block_height=0.02, block_width = 0.02,
        viewer_image_height = 84, viewer_image_width = 84, skip_first=100,
        substeps=100, init_hand_xyz=(0, 0.4, 0.07),
        randomize_initial_pos = False, state_goal = None,
        randomize_goal_at_reset = False, fix_rotation = True,
        fix_reset_pos = True, fix_gripper = True, do_render = False,
        match_orientation = False,
        workspace_low = np.array([-0.2, 0.4, 0.05]),
        workspace_high = np.array([0.2, 0.8, 0.1]),
        hand_low = np.array([-0.2, 0.4, 0.05]),
        hand_high = np.array([0.2, 0.8, 0.1]),
    ):
        self.quick_init(locals())

        base_filename = asset_base_path + filename

        friction_params = (friction, 0.1, 0.02)
        self.obj_stat_prop = create_object_xml(base_filename, num_objects, object_mass,
                                               friction_params, object_meshes, finger_sensors,
                                               maxlen, minlen, preload_obj_dict, obj_classname,
                                               block_height, block_width)


        gen_xml = create_root_xml(base_filename)
        SawyerXYZEnv.__init__(self, model_name=gen_xml)
        clean_xml(gen_xml)

        if self.sim.model.nmocap > 0 and self.sim.model.eq_data is not None:
            for i in range(self.sim.model.eq_data.shape[0]):
                if self.sim.model.eq_type[i] == mujoco_py.const.EQ_WELD:
                    # Define the xyz + quat of the mocap relative to the hand
                    self.sim.model.eq_data[i, :] = np.array(
                        [0., 0., 0., 1., 0., 0., 0.]
                    )

        self._base_sdim, self._base_adim, self.mode_rel = 5, 5, mode_rel
        self.hand_low = hand_low
        self.hand_high = hand_high
        self.mocap_high = self.hand_high
        self.mocap_low = self.hand_low
        self.action_scale = 1.0/100
        self.num_objects, self.skip_first, self.substeps = num_objects, skip_first, substeps
        self.randomize_initial_pos = randomize_initial_pos
        self.finger_sensors, self._maxlen = finger_sensors, maxlen
        self._threshold = 0
        self._previous_target_qpos, self._n_joints = None, 8
        self._object_pos = np.zeros((num_objects, 7))
        self._reset_xyz = np.zeros((num_objects, 3))
        self._reset_quat = np.zeros((num_objects, 4))
        self._reset_pos = np.copy(self._object_pos)
        self._object_names = ['obj_' + str(i) for i in range(num_objects)]
        self._initialized = False
        self._state_goal = state_goal
        self.fix_rotation = fix_rotation
        self.fix_gripper = fix_gripper
        self._randomize_goal_at_reset = randomize_goal_at_reset
        self.do_render = do_render
        self.init_hand_xyz = np.array(init_hand_xyz)
        self.match_orientation = match_orientation
        self.goal_dim_per_object = 7 if match_orientation else 3
        self.workspace_low = workspace_low
        self.workspace_high = workspace_high

        # self.finger_sensors = False # turning this off for now
        self.action_space = Box(np.array([-1, -1, -1, -1, -1]),
                        np.array([1, 1, 1, 1, 1]), dtype=np.float32)
        # action space is 5-dim: x, y, z, ? ?
        self.hand_space = Box(self.hand_low, self.hand_high, dtype=np.float32)
        obs_dim = 7 * num_objects + 3
        obs_space = Box(np.array([-1] * obs_dim), np.array([1] * obs_dim), dtype=np.float32)
        goal_space = Box(
            np.array([-1] * self.goal_dim_per_object * num_objects),
            np.array([1] * self.goal_dim_per_object * num_objects),
            dtype=np.float32,
        )
        self.observation_space = Dict([
            ('observation', obs_space),
            ('desired_goal', goal_space),
            ('achieved_goal', goal_space),
            ('state_observation', obs_space),
            ('state_desired_goal', goal_space),
            ('state_achieved_goal', goal_space),
        ])

        self._state_goal = self.sample_goal()
        self.fix_reset_pos = fix_reset_pos
        o = self.reset()

    def _clip_gripper(self):
        self.sim.data.qpos[7:9] = np.clip(self.sim.data.qpos[7:9], [-0.055, 0.0027], [-0.0027, 0.055])

    def samp_xyz_rot(self):
        rand_xyz = np.random.uniform(
            self.workspace_low[:3] + self._maxlen / 2 + 0.02,
            self.workspace_high[:3] - self._maxlen / 2 + 0.02
        )
        rand_xyz[-1] = 0.05
        return rand_xyz, np.random.uniform(-np.pi / 2, np.pi / 2)

    def project_point(self, point, camera):
        model_matrix = np.zeros((4, 4))
        model_matrix[:3, :3] = self.sim.data.get_camera_xmat(camera).T
        model_matrix[-1, -1] = 1

        fovy_radians = np.deg2rad(self.sim.model.cam_fovy[self.sim.model.camera_name2id(camera)])
        uh = 1. / np.tan(fovy_radians / 2)
        uw = uh / (self._frame_width / self._frame_height)
        extent = self.sim.model.stat.extent
        far, near = self.sim.model.vis.map.zfar * extent, self.sim.model.vis.map.znear * extent
        view_matrix = np.array([[uw, 0., 0., 0.],                        #matrix definition from
                                [0., uh, 0., 0.],                        #https://stackoverflow.com/questions/18404890/how-to-build-perspective-projection-matrix-no-api
                                [0., 0., far / (far - near), -1.],
                                [0., 0., -2*far*near/(far - near), 0.]]) #Note Mujoco doubles this quantity

        MVP_matrix = view_matrix.dot(model_matrix)
        world_coord = np.ones((4, 1))
        world_coord[:3, 0] = point - self.sim.data.get_camera_xpos(camera)

        clip = MVP_matrix.dot(world_coord)
        ndc = clip[:3] / clip[3]  # everything should now be in -1 to 1!!
        col, row = (ndc[0] + 1) * self._frame_width / 2, (-ndc[1] + 1) * self._frame_height / 2

        return self._frame_height - row, col                 #rendering flipped around in height

    def get_desig_pix(self, cams, target_width, round=True):
        qpos_dim = self._n_joints      # the states contains pos and vel
        assert self.sim.data.qpos.shape[0] == qpos_dim + 7 * self.num_objects
{'fraction_non_alphanumeric': 0.12120844455229313, 'fraction_numerical': 0.03652026207231254, 'mean_word_length': 2.4749156829679597, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 9, 'https://': 4, 'lorem ipsum': 0, 'www.': 0, 'xml': 22}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 14, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1085955', 'n_tokens_mistral': 6104, 'n_tokens_neox': 5603, 'n_words': 1143}