intuitive-robots · xinkai-jiang · Nov 17, 2024 · Nov 12, 2024 · Nov 15, 2024 · Nov 15, 2024
diff --git a/.gitignore b/.gitignore
@@ -6,7 +6,7 @@ __pycache__/
 scenes/
 models copy/
 test.py
-*.xml
+# *.xml
 # C extensions
 *.so
 *.json

diff --git a/demos/fancy_gym/TableTennis.py b/demos/fancy_gym/TableTennis.py
@@ -8,7 +8,7 @@
 env = fancy_gym.make(env_name, seed=1)
 obs = env.reset()
 
-publisher = FancyGymPublisher(env, "127.0.0.1")
+publisher = FancyGymPublisher(env, "192.168.0.134")
 
 while True:
     action = env.action_space.sample()

diff --git a/demos/libero.py b/demos/libero.py
@@ -0,0 +1,67 @@
+import numpy as np
+
+from robosuite.environments.base import MujocoEnv as RobosuiteEnv
+import xml.etree.ElementTree as ET
+import os
+
+from simpub.xr_device.meta_quest3 import MetaQuest3 as MetaQuest3Sim
+from simpub.sim.mj_publisher import MujocoPublisher
+
+
+from robosuite import load_controller_config
+import libero.libero.envs.bddl_utils as BDDLUtils
+from libero.libero.envs import TASK_MAPPING
+from robosuite.robots import ROBOT_CLASS_MAPPING
+
+
+class RobosuitePublisher(MujocoPublisher):
+
+    def __init__(self, env: RobosuiteEnv):
+        super().__init__(
+            env.sim.model._model,
+            env.sim.data._data,
+            visible_geoms_groups=[1, 2, 3, 4]
+        )
+
+
+if __name__ == "__main__":
+    # Get controller config
+    controller_config = load_controller_config(default_controller="OSC_POSE")
+
+    # Create argument configuration
+    config = {
+        "robots": ["Panda"],
+        "controller_configs": controller_config,
+    }
+
+    bddl_path = "/home//LIBERO/libero/libero/bddl_files/"
+    bddl_dataset_name = "libero_10"
+    bddl_name = "LIVING_ROOM_SCENE2_put_both_the_cream_cheese_box_and_the_butter_in_the_basket.bddl"
+    bddl_file = os.path.join(bddl_path, bddl_dataset_name, bddl_name)
+    problem_info = BDDLUtils.get_problem_info(bddl_file)
+    # Create environment
+    problem_name = problem_info["problem_name"]
+    domain_name = problem_info["domain_name"]
+    language_instruction = problem_info["language_instruction"]
+    if "TwoArm" in problem_name:
+        config["env_configuration"] = "single-arm-opposed"
+    print(language_instruction)
+    env = TASK_MAPPING[problem_name](
+        bddl_file_name=bddl_file,
+        **config,
+        has_renderer=True,
+        has_offscreen_renderer=False,
+        render_camera="agentview",
+        ignore_done=True,
+        use_camera_obs=False,
+        reward_shaping=True,
+        control_freq=20,
+    )
+
+    # reset the environment
+    env.reset()
+    publisher = RobosuitePublisher(env)
+    while True:
+        action = np.random.randn(env.robots[0].dof - 1)  # sample random action
+        obs, reward, done, info = env.step(action)  # take action in the environment
+        env.render()  # render on display
diff --git a/demos/mujoco/table_tennis_two_players/assets/bats.xml b/demos/mujoco/table_tennis_two_players/assets/bats.xml
@@ -0,0 +1,12 @@
+<mujocoinclude>
+    <body name="bat1" pos="1 0 1" childclass="contact_geom">
+        <geom name="bat1" type="cylinder" size="0.075 0.005" rgba="1 0 0 1" pos="-0.05 0 0" quat="0 0 1 0"/>
+        <!-- <geom name="bat1_back" type="cylinder" size="0.075 0.0025" rgba="0 1 0 1" quat="0 0 1 0" pos="0 -0.0026 0"/> -->
+        <geom name="handle1" type="cylinder" size="0.015 0.025" rgba="0.6 0.3 0 1" pos="0 0 0" quat="0.71 0 0.71 0"/>
+    </body>
+    <body name="bat2" pos="-1 0 1" childclass="contact_geom">
+        <geom name="bat2" type="cylinder" size="0.075 0.005" rgba="1 0 0 1" quat="0.71 0 0.71 0"/>
+        <geom name="bat2_back" type="cylinder" size="0.075 0.0025" rgba="0 1 0 1" quat="0.71 0 0.71 0" pos="-0.0026 0 0"/>
+        <geom name="handle2" type="cylinder" size="0.015 0.025" rgba="0.6 0.3 0 1" pos="0 0 0.05" quat="1 0 0 0"/>
+    </body>
+</mujocoinclude>
diff --git a/demos/mujoco/table_tennis_two_players/assets/include_free_ball.xml b/demos/mujoco/table_tennis_two_players/assets/include_free_ball.xml
@@ -0,0 +1,11 @@
+<mujocoinclude>
+    <body name="target_ball" pos="0 0 0">
+<!--        <joint axis="1 0 0" damping="0.0" name="tar:x" pos="0 0 0" stiffness="0" type="slide" frictionloss="0"/>-->
+<!--        <joint axis="0 1 0" damping="0.0" name="tar:y" pos="0 0 0" stiffness="0" type="slide" frictionloss="0"/>-->
+<!--        <joint axis="0 0 1" damping="0.0" name="tar:z" pos="0 0 0" stiffness="0" type="slide" frictionloss="0"/>-->
+        <joint name="ball_joint" damping="0.0" pos="0 0 0" stiffness="0" frictionloss="0" type="free"/>
+        <geom size="0.025 0.025 0.025" type="sphere" condim="4" name="target_ball_contact" rgba="1 1 0 1" mass="0.1"
+              friction="0.1 0.1 0.1" solimp="0.9 0.95 0.001 0.5 2" solref="0.1 0.03" priority="1"/>
+        <site name="target_ball" pos="0 0 0" size="0.02 0.02 0.02" rgba="1 0 0 1" type="sphere"/>
+    </body>
+</mujocoinclude>
diff --git a/demos/mujoco/table_tennis_two_players/assets/include_table.xml b/demos/mujoco/table_tennis_two_players/assets/include_table.xml
@@ -0,0 +1,30 @@
+<mujocoinclude>
+    <body name="table_tennis_table" pos="0 0 0">
+        <geom class="viz" name="table_base_1" pos="1 0.7 0.375" rgba="1 1 1 1" size="0.05 0.05 .375" type="box" />
+        <geom class="viz" name="table_base_2" pos="1 -0.7 0.375" rgba="1 1 1 1" size="0.05 0.05 .375" type="box" />
+        <geom class="viz" name="table_base_3" pos="-1 -0.7 0.375" rgba="1 1 1 1" size="0.05 0.05 .375" type="box" />
+        <geom class="viz" name="table_base_4" pos="-1 0.7 0.375" rgba="1 1 1 1" size="0.05 0.05 .375" type="box" />
+        <body name="table_top" pos="0 0 0.76">
+            <geom class="contact_geom" name="table_tennis_table" pos="0 0 0" rgba="0 0 0.5 1" size="1.37 .7625 .01" type="box" />
+            <site name="left_up_corner" pos="-1.37 .7625 0.01" rgba="1 1 1 1" size="0.00001" />
+            <site name="middle_up_corner" pos="0 .7625 0.01" rgba="1 1 1 1" size="0.00001" />
+            <site name="left_down_corner" pos="-1.37 -0.7625 0.01" rgba="1 1 1 1" size="0.00001" />
+            <site name="middle_down_corner" pos="0 -.7625 0.01" rgba="1 1 1 1" size="0.00001" />
+            <geom class="contact_geom" material="floor_plane" name="table_te_context_spacennis_net" pos="0 0 0.08625" rgba="0 0 1 0.5" size="0.01 0.915 0.07625" type="box" />
+            <geom class="viz" name="left_while_line" pos="0 -0.7425 0.01" rgba="1 1 1 1" size="1.37 .02 .001" type="box" />
+            <geom class="viz" name="center_while_line" pos="0 0 0.01" rgba="1 1 1 1" size="1.37 .01 .001" type="box" />
+            <geom class="viz" name="right_while_line" pos="0 0.7425 0.01" rgba="1 1 1 1" size="1.37 .02 .001" type="box" />
+            <geom class="viz" name="right_side_line" pos="1.35 0 0.01" rgba="1 1 1 1" size="0.02 .7625 .001" type="box" />
+            <geom class="viz" name="left_side_line" pos="-1.35 0 0.01" rgba="1 1 1 1" size="0.02 .7625 .001" type="box" />
+        </body>
+        <body name="achieved_pos" pos="0 0 0.5">
+            <geom class="viz" name="achieved_point_geom" pos="0 0 0" rgba="0 1 0 1" size="0.02 0.001" type="cylinder" />
+        </body>
+        <body name="right_achieved_pos" pos="0 0 0.5">
+            <geom class="viz" name="hitting_achieved_point_geom" pos="0 0 0" rgba="1 0 0 1" size="0.02 0.001" type="cylinder" />
+        </body>
+        <body name="target_point" pos="0 0 0.5">
+            <geom class="viz" name="target_point_geom" pos="0 0 0" rgba="1 1 0 1" size="0.02 0.001" type="cylinder" />
+        </body>
+    </body>
+</mujocoinclude>
diff --git a/demos/mujoco/table_tennis_two_players/assets/include_target_ball.xml b/demos/mujoco/table_tennis_two_players/assets/include_target_ball.xml
@@ -0,0 +1,10 @@
+<mujocoinclude>
+    <body name="target_ball" pos="0. 0. 2">
+        <joint axis="1 0 0" damping="0.0" name="tar_x" pos="0 0 0" stiffness="0" type="slide" frictionloss="0"/>
+        <joint axis="0 1 0" damping="0.0" name="tar_y" pos="0 0 0" stiffness="0" type="slide" frictionloss="0"/>
+        <joint axis="0 0 1" damping="0.0" name="tar_z" pos="0 0 0" stiffness="0" type="slide" frictionloss="0"/>
+        <geom size="0.025 0.025 0.025" type="sphere" condim="4" name="target_ball_contact" rgba="1 1 0 1" mass="0.1"
+              friction="0.1 0.1 0.1" solimp="0.9 0.95 0.001 0.5 2" solref="0.1 0.03" priority="1"/>
+        <site name="target_ball" pos="0 0 0" size="0.02 0.02 0.02" rgba="1 0 0 1" type="sphere"/>
+    </body>
+</mujocoinclude>
diff --git a/demos/mujoco/table_tennis_two_players/assets/table_tennis_env.xml b/demos/mujoco/table_tennis_two_players/assets/table_tennis_env.xml
@@ -0,0 +1,28 @@
+<mujoco model="table_tennis(v0.1)">
+    <compiler angle="radian" coordinate="local" />
+    <option gravity="0 0 -9.81" timestep="0.002">
+        <flag warmstart="enable" />
+    </option>
+    <custom>
+        <numeric data="0 0 0 0 0 0 0" name="START_ANGLES" />
+    </custom>
+    <default>
+        <default class="viz">
+            <geom type="mesh" contype="0" conaffinity="0" group="1" rgba="1 1 1 1"/>
+        </default>
+        <default class="contact_geom">
+            <geom condim="4" friction="0.1 0.1 0.1" margin="0" solimp="1 1 0" solref="0.1 0.03"/>
+        </default>
+    </default>
+    <asset>
+        <texture builtin="checker" height="512" name="texplane" rgb1=".2 .3 .4" rgb2=".1 0.15 0.2" type="2d"
+                 width="512"/>
+        <material name="floor_plane" reflectance="0.5" texrepeat="1 1" texture="texplane" texuniform="true"/>
+    </asset>
+    <worldbody>
+        <light cutoff="60" diffuse="1 1 1" dir="-.1 -.2 -1.3" directional="true" exponent="1" pos=".1 .2 1.3" specular=".1 .1 .1" />
+        <include file="include_table.xml" />
+        <include file="include_target_ball.xml" />
+        <include file="bats.xml" />
+    </worldbody>
+</mujoco>
diff --git a/demos/mujoco/table_tennis_two_players/table_tennis.py b/demos/mujoco/table_tennis_two_players/table_tennis.py
@@ -0,0 +1,53 @@
+import mujoco
+from mujoco import mj_name2id, mjtObj
+import numpy as np
+import time
+from simpub.sim.mj_publisher import MujocoPublisher
+from simpub.xr_device.meta_quest3 import MetaQuest3
+
+
+def check_episode_and_rest(mj_model, mj_data):
+    target_ball_id = mj_name2id(model, mjtObj.mjOBJ_BODY, "target_ball")
+    object_position = mj_data.xpos[target_ball_id]
+    if abs(object_position[0]) > 2 or abs(object_position[1]) > 2:
+        body_jnt_addr = mj_model.body_jntadr[target_ball_id]
+        qposadr = mj_model.jnt_qposadr[body_jnt_addr]
+        mj_data.qpos[qposadr:qposadr + 3] = np.array([0, 0, 2])
+        mj_data.qvel[qposadr:qposadr + 3] = np.array([0, 0, 0])
+
+
+def update_bat(mj_model, mj_data, player1: MetaQuest3, player2: MetaQuest3 = None):
+    bat1_id = mj_name2id(model, mjtObj.mjOBJ_BODY, "bat1")
+    player1_input = player1.get_input_data()
+    # print(player1_input)
+    if player1_input is None:
+        return
+    mj_model.body_pos[bat1_id] = np.array(player1_input["right"]["pos"])
+    quat = player1_input["right"]["rot"]
+    mj_model.body_quat[bat1_id] = np.array([quat[3], quat[0], quat[1], quat[2]])
+    # bat2_id = mj_name2id(model, mjtObj.mjOBJ_BODY, "bat2")
+    # mj_data.body_pos[bat2_id] = np.array(player1.input_data["left"]["pos"])
+    # mj_data.body_quat[bat2_id] = np.array(player1.input_data["left"]["rot"])
+
+
+if __name__ == '__main__':
+
+    model = mujoco.MjModel.from_xml_path("assets/table_tennis_env.xml")
+    data = mujoco.MjData(model)
+    last_time = time.time()
+    publisher = MujocoPublisher(model, data, host="192.168.0.134")
+    player1 = MetaQuest3("UnityClient")
+    # player2 = MetaQuest3("ALR2")
+    while not player1.connected:
+        time.sleep(0.01)
+    print("Connected to UnityClient")
+    count = 0
+    while True:
+        mujoco.mj_step(model, data)
+        if time.time() - last_time < 0.001:
+            time.sleep(0.001 - (time.time() - last_time))
+        last_time = time.time()
+        if count % 10 == 0:
+            update_bat(model, data, player1)
+            check_episode_and_rest(model, data)
+        count += 1
diff --git a/demos/robocasa/kitchen.py b/demos/robocasa/kitchen.py
@@ -13,6 +13,7 @@
 from robocasa.scripts.collect_demos import collect_human_trajectory
 
 from simpub.sim.robocasa_publisher import RobocasaPublisher
+from simpub.xr_device.meta_quest3 import MetaQuest3
 
 if __name__ == "__main__":
     # Arguments
@@ -78,7 +79,8 @@
     # return None, True
     env.render()
 
-    publisher = RobocasaPublisher(env)
+    publisher = RobocasaPublisher(env, host="192.168.0.134")
+    meta = MetaQuest3("UnityClient")
 
     while True:
         zero_action = np.zeros(env.action_dim)

diff --git a/simpub/parser/mj.py b/simpub/parser/mj.py
@@ -25,9 +25,9 @@ def sphere2unity_scale(scale: List[float]) -> List[float]:
 
 def cylinder2unity_scale(scale: List[float]) -> List[float]:
     if len(scale) == 3:
-        return list(map(abs, [scale[0], scale[1], scale[0]]))
+        return list(map(abs, [scale[0] * 2, scale[1], scale[0] * 2]))
     if len(scale) == 2:
-        return list(map(abs, [scale[0], scale[1], scale[0]]))
+        return list(map(abs, [scale[0] * 2, scale[1], scale[0] * 2]))
     elif len(scale) == 1:
         return list(map(abs, [scale[0] * 2, scale[0] * 2, scale[0] * 2]))
 
@@ -133,19 +133,11 @@ def process_body(self, mj_model, body_id: int):
             mj_model.body_geomadr[body_id],
             mj_model.body_geomadr[body_id] + num_geoms
         ):
-            geom_name = mujoco.mj_id2name(
-                mj_model, mujoco.mjtObj.mjOBJ_GEOM, geom_id
-            )
             geom_group = int(mj_model.geom_group[geom_id])
             # check if the geom participates in rendering
             if geom_group not in self.visible_geoms_groups:
-                logger.info(
-                    (
-                        f"Geom '{geom_name}'(id {geom_id}) does not"
-                        f"participate in rendering and can be removed."
-                    )
-                )
                 continue
+
             sim_visual = self.process_geoms(mj_model, geom_id)
             sim_object.visuals.append(sim_visual)
         return sim_object, parent_id
@@ -193,7 +185,7 @@ def process_mesh(self, mj_model, mesh_id: int):
         vertices = mj_model.mesh_vert[start_vert:start_vert + num_verts]
         vertices = vertices.astype(np.float32)
         vertices = vertices[:, [1, 2, 0]]
-        vertices[:, 0] = - vertices[:, 0]
+        vertices[:, 0] = -vertices[:, 0]
         vertices = vertices.flatten()
         vertices_layout = bin_buffer.tell(), vertices.shape[0]
         bin_buffer.write(vertices)
@@ -207,7 +199,7 @@ def process_mesh(self, mj_model, mesh_id: int):
         norms = mj_model.mesh_normal[start_norm:start_norm + num_norm]
         norms = norms.astype(np.float32)
         norms = norms[:, [1, 2, 0]]
-        norms[:, 0] = - norms[:, 0]
+        norms[:, 0] = -norms[:, 0]
         norms = norms.flatten()
         normal_layout = bin_buffer.tell(), norms.shape[0]
         bin_buffer.write(norms)
@@ -289,8 +281,8 @@ def process_texture(self, mj_model, tex_id: int):
         # TODO: Texture type?
         # tex_type = mj_model.tex_type[tex_id]
         # get the texture data
-        tex_height = mj_model.tex_height[tex_id]
-        tex_width = mj_model.tex_width[tex_id]
+        tex_height = mj_model.tex_height[tex_id].item()
+        tex_width = mj_model.tex_width[tex_id].item()
         # only we only supported texture channel number is 3
         if hasattr(mj_model, "tex_nchannel"):
             tex_nchannel = mj_model.tex_nchannel[tex_id]
@@ -309,18 +301,25 @@ def process_texture(self, mj_model, tex_id: int):
             ]
 
         # compress the texture data
-        width = int(tex_width // 4)
-        height = int(tex_height // 4)
-        tex_data = cv2.resize(
-            tex_data.reshape(tex_width, tex_height, 3),
-            (width, height),
-            interpolation=cv2.INTER_LINEAR
-        )
+        max_texture_size = 10 ** 6 / 2 # ~ 500 kB (robocasa scene 50 images -> 25 MB)
+
+        scale = np.sqrt(len(tex_data) / max_texture_size)  
+        if scale > 1:
+            scale = int(scale) + 1
+
+            new_width, new_height = tex_width // scale, tex_height // scale
+            tex_data = cv2.resize(
+                tex_data.reshape(tex_width, tex_height, -1),
+                (new_width, new_height),
+                interpolation=cv2.INTER_LINEAR
+            )
+            tex_width, tex_height = new_height, new_height
+
         bin_data = tex_data.astype(np.uint8).tobytes()
         texture_hash = md5(bin_data).hexdigest()
         texture = SimTexture(
-            width=width,
-            height=height,
+            width=tex_width,
+            height=tex_height,
             # Only support 2d texture
             textureType="2D",
             hash=texture_hash

diff --git a/simpub/xr_device/xr_device.py b/simpub/xr_device/xr_device.py
@@ -35,12 +35,12 @@ def __init__(
     async def wait_for_connection(self):
         logger.info(f"Waiting for connection to {self.device_name}")
         while not self.connected:
-            client = self.manager.clients.get(self.device_name)
-            if client is not None:
-                self.connected = True
-                self.client = client
-                logger.info(f"Connected to {self.device_name}")
-                break
+            for client in self.manager.clients.values():
+                if client.info["name"] == self.device_name:
+                    self.client = client
+                    self.connected = True
+                    logger.info(f"Connected to {self.device_name}")
+                    break
             await asycnc_sleep(0.01)
         self.req_socket = self.client.req_socket
         self.sub_socket = self.client.sub_socket
-Original file line number
+Diff line change
@@ Expand Up / @@ -6,7 +6,7 @@ __pycache__/ @@
     scenes/
     models copy/
     test.py
-    *.xml
+    # *.xml
     # C extensions
     *.so
     *.json
@@ Expand Down @@