Gazebo cleanup merge

.gitignore

@@ -33,4 +33,8 @@ __pycache__
 
 # No wandb outputs or runs
 **/wandb/
-**/runs/
+**/runs/
+
+mirage/mirage/ros_ws/build/*
+mirage/mirage/ros_ws/install/*
+mirage/mirage/ros_ws/log/*

README.md

@@ -22,31 +22,66 @@ Clone the repo:
 git clone --recurse-submodules [email protected]:BerkeleyAutomation/mirage.git
 ```
 
-Create a conda environment or virtualenv:
-```
-conda create --name mirage python=3.8
-```
+Install ROS 2 and setup the Gazebo Environment for Inpainting by following the instructions in `mirage/mirage/ros_ws/src/README.md` 
 
-Install the mirage Python package.
+Create a conda environment and install the mirage Python package.
 ```
+conda create -n mirage --python=3.10
 cd mirage
 pip install -e .
 ```
 
-Follow the installation instructions for the given simulator by going into the forked repositories.
+Follow the installation instructions for the downloaded robomimic, robosuite, and mimicgen_environments submodules by going into them and looking at the README.
+
+Build the ROS workspace by running (from the root of this repo)
+```
+cd mirage/mirage/ros_ws
+colcon build
+source install/setup.bash
+```
 
 ## Usage
-For robosuite, to run an experiment, 
+### Robosuite Benchmark
+If inpainting is enabled, only 1 benchmarking process can be run at a given time.
+Depending on the robosuite environment, the ros launch file will be different due to different camera extrinsics.
+Firstly, in one terminal, run the gazebo writer node
+```
+source mirage/mirage/ros_ws/install/setup.bash
+ros2 run gazebo_env write_data_node_robosuite_better.py
+```
+
+Then, launch the gazebo process for the corresponding environment by first running (in a new terminal):
+```
+source mirage/mirage/ros_ws/install/setup.bash
+```
+Run either one of these depending on the environment (three piece assembly is really two piece assembly).
+```
+ros2 launch gazebo_env panda_gazebo_classic_robosuite_can.launch.py
+ros2 launch gazebo_env panda_gazebo_classic_robosuite_lift_square_stack_three_threading.launch.py
+ros2 launch gazebo_env panda_gazebo_classic_robosuite_three_piece_assembly.launch.py
 ```
-cd mirage/benchmark/robosuite
-python3 run_robosuite_benchmark.py --config_file config/example_config.yaml
+
+For robosuite, to run an experiment (from the root of this repo), 
+```
+# Run this command below if inpainting, otherwise skip
+source mirage/mirage/ros_ws/install/setup.bash
+
+cd mirage/mirage/benchmark/robosuite
+python3 run_robosuite_benchmark.py --config config/example_config.yaml
 ```
 Please take a look at the example_config and the different parameters that can be set to run different tasks, agents, and robots. For the above code to work, you must change the agents to the path for the model checkpoints in robosuite.
 
 ## Citation
 If you utilized the benchmark, please consider citing the paper:
 ```
-TODO: Add link to ArXiv / publication
+@misc{chen2024mirage,
+      title={Mirage: Cross-Embodiment Zero-Shot Policy Transfer with Cross-Painting}, 
+      author={Lawrence Yunliang Chen and Kush Hari and Karthik Dharmarajan and Chenfeng Xu and Quan Vuong and Ken Goldberg},
+      year={2024},
+      eprint={2402.19249},
+      archivePrefix={arXiv},
+      primaryClass={cs.RO}
+}
 ```
 
 ## Contributing

mirage/mirage/benchmark/robosuite/evaluate_policy_demo_source_robot_server.py

@@ -155,6 +155,7 @@ def __init__(self, robot_name=None, ckpt_path=None, render=False, video_path=Non
                 self.rollout_horizon = config.experiment.rollout.horizon
 
         self.core_env = self.env.env
+
         self.is_diffusion = False
         while hasattr(self.core_env, "env"):
             self.is_diffusion = True

mirage/mirage/benchmark/robosuite/evaluate_policy_demo_target_robot_client.py

@@ -52,6 +52,17 @@ def __init__(self, robot_name=None, ckpt_path=None, render=False, video_path=Non
                         else:
                             print("NOT USING ROS! Will use the groundtruth Franka masks instead")                        
 
+        self.should_kill_early = False
+
+        import threading
+        self.thread = threading.Thread(target=self.wait_for_skip)
+        self.thread.start()
+
+    def wait_for_skip(self):
+        while True:
+            input()
+            self.should_kill_early = True
+
     def image_to_pointcloud(self, depth_map, camera_name, camera_height=84, camera_width=84, segmask=None):
         """
         Convert depth image to point cloud
@@ -205,6 +216,14 @@ def rollout_robot(self, video_skip=5, return_obs=False, camera_names=None, set_o
                             from mirage.infra.ros_inpaint_publisher_sim import ROSInpaintSimData
                             eef_pose = self.compute_eef_pose()
                             eef_pose_matrix = T.pose2mat((eef_pose[:3], eef_pose[3:]))
+
+                            # For some reason the Can / Pick and Place task is very weird
+                            # It has a y offset of -0.1 for the base of the robot, so there needs to be a
+                            # 0.1 offset for the ee to match ground truth, this needs more investigation....
+                            from robosuite.environments.manipulation.pick_place import PickPlace
+                            if isinstance(self.core_env, PickPlace):
+                                eef_pose_matrix[1] += 0.1
+
                             data = ROSInpaintSimData(ros_rgb_img, ros_depth_img, ros_segmentation_mask, eef_pose_matrix, obs['robot0_gripper_qpos'][-1:])
                             print("Joints including gripper", sent_joint_angles)
                             self.ros_inpaint_publisher.publish_to_ros_node(data)
@@ -340,6 +359,11 @@ def rollout_robot(self, video_skip=5, return_obs=False, camera_names=None, set_o
                 if step_i - source_finished_step >= 10:
                     done = True
 
+            # If human hits enter, kill off the run, it isn't feasible
+            if self.should_kill_early:
+                self.should_kill_early = False
+                done = True
+
             # tell the source robot that the target robot is ready
             # Create an instance of Data() to send to client.
             variable = Data()

mirage/mirage/benchmark/robosuite/robosuite_experiment.py

@@ -98,7 +98,6 @@ def launch(self, override=False) -> None:
             target_agent_args.append("--inpaint_enabled")
 
             if self._config.use_ros:
-                source_agent_args.append("--use_ros")
                 target_agent_args.append("--use_ros")
 
             if self._config.offline_eval:

mirage/mirage/gripper_interpolation/robosuite/gripper_interpolator.py

@@ -50,8 +50,12 @@ def interpolate_gripper(self, gripper_angles: np.array) -> np.array:
         """
         if self.source_robot == self.target_robot:
             print("Returning angles")
-            return 1 * gripper_angles
+            return -1 * gripper_angles
 
+        if (self.source_robot == "Panda" or self.source_robot == "Kinova3") and self.target_robot == "UR5e": 
+        # TODO: add this for real gripper mapping
+            return -0.05702400673569841 * gripper_angles + 0.02670973458948458
+
         relevant_interpoolator = self.interpolators[(self.source_robot, self.target_robot)]
         regression_model = LinearRegression()
 

mirage/mirage/infra/ros_inpaint_publisher_sim.py

@@ -1,6 +1,6 @@
 from input_filenames_msg.msg import InputFilesSimData
 from mirage.infra.ros_inpaint_publisher import ROSInpaintPublisher
-from mirage.mirage.gripper_interpolation.gripper_interpolator import GripperInterpolator
+from mirage.gripper_interpolation.robosuite.gripper_interpolator import GripperInterpolator
 import numpy as np
 
 class ROSInpaintSimData:
@@ -30,7 +30,7 @@ class ROSInpaintPublisherSim(ROSInpaintPublisher):
     to a node that performs inpainting on a target robot.
     """
 
-    def __init__(self, source_robot_info: str, target_robot_info: str):
+    def __init__(self):
         """
         Initializes the ROS2 node.
         :param source_robot_info: the information about the source robot to determine which interpolation scheme to use
@@ -41,11 +41,8 @@ def __init__(self, source_robot_info: str, target_robot_info: str):
         self._publisher = self.node.create_publisher(
             InputFilesSimData, '/input_files_data_sim', 1)
 
-        # TODO: generalize this
-        self.gripper_interpolator = GripperInterpolator(source_robot_info, target_robot_info, ['/home/mirage/x-embody/xembody/xembody_robosuite/paired_trajectories_collection/gripper_interpolation_results_no_task_diff.pkl',
-                                                                                               '/home/mirage/x-embody/xembody/xembody_robosuite/paired_trajectories_collection/gripper_interpolation_results_20_rollouts.pkl'])
-        # self.gripper_interpolator = GripperInterpolator('panda', 'panda', ['/home/mirage/x-embody/xembody/xembody_robosuite/paired_trajectories_collection/gripper_interpolation_results_no_task_diff.pkl'])
-        # self.gripper_interpolator = GripperInterpolator('panda', 'ur5', '/home/mirage/x-embody/xembody/xembody_robosuite/paired_trajectories_collection/gripper_interpolation_results_no_task_diff.pkl')
+        # TODO(kdharmarajandev): generalize this
+        self.gripper_interpolator = GripperInterpolator('Panda', 'Panda')
 
     def publish_to_ros_node(self, data: ROSInpaintSimData):
         """

mirage/mirage/ros_ws/src/README.md

@@ -1,8 +1,81 @@
-# Cross Embodiment
+# Mirage
 
-Hi. Welcome to this test workspace for the cross embodiment project. This is currently in development.
+Hi! Welcome to the ROS workspace setup for Mirage. This README details how to setup the ROS2 environment and Gazebo to run simulated and real experiments in the Mirage paper (RSS 2024 hopefully ;)
 
-## Installation
+## Prerequisites
+Computer with Ubuntu 22.04
+
+## Installation (Bash Script)
+Clone the repo: 
+```
+cd ~/
+git clone --recurse-submodules [email protected]:BerkeleyAutomation/mirage.git
+```
+
+Run the setup bash script
+```
+cd ~/mirage/xembody/xembody/src/ros_ws/src
+bash ros_gazebo_setup.bash
+```
+
+## Installation (Step by Step)
+
+### Step 1: Install ROS2 Humble Full: https://docs.ros.org/en/humble/Installation/Ubuntu-Install-Debians.html
+```
+# Step 1.1 Set Locale
+sudo apt update -y
+sudo apt install locales
+sudo locale-gen en_US en_US.UTF-8
+sudo update-locale LC_ALL=en_US.UTF-8 LANG=en_US.UTF-8
+export LANG=en_US.UTF-8
+
+# Step 1.2 Setup Sources
+sudo apt install software-properties-common
+echo -ne '\n' | sudo add-apt-repository universe
+sudo apt update && sudo apt install curl -y
+sudo curl -sSL https://raw.githubusercontent.com/ros/rosdistro/master/ros.key -o /usr/share/keyrings/ros-archive-keyring.gpg
+echo "deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/ros-archive-keyring.gpg] http://packages.ros.org/ros2/ubuntu $(. /etc/os-release && echo $UBUNTU_CODENAME) main" | sudo tee /etc/apt/sources.list.d/ros2.list > /dev/null
+
+# Step 1.3 Install ROS2 Packages
+sudo apt update -y
+sudo apt upgrade -y
+sudo apt install ros-humble-desktop-full -y
+echo "source /opt/ros/humble/setup.bash" >> ~/.bashrc
+```
+
+### Step 2: Setup ROS2 Control: https://control.ros.org/humble/doc/ros2_control_demos/doc/index.html
+```
+mkdir -p ~/ros2_ws/src
+cd ~/ros2_ws/src
+git clone https://github.com/ros-controls/ros2_control_demos -b humble
+cd ~/ros2_ws/
+sudo apt-get update
+sudo rosdep install --from-paths ./ -i -y --rosdistro ${ROS_DISTRO}
+cd ~/ros2_ws/
+. /opt/ros/${ROS_DISTRO}/setup.sh
+colcon build --merge-install
+```
+
+### Step 3: Tracikpy Installation: https://github.com/mjd3/tracikpy
+```
+sudo apt-get install libeigen3-dev liborocos-kdl-dev libkdl-parser-dev liburdfdom-dev libnlopt-dev
+cd ~/
+git clone https://github.com/mjd3/tracikpy.git
+pip install tracikpy/
+```
+
+### Step 4: Make sure Gazebo plugins can be found: https://classic.gazebosim.org/tutorials?tut=guided_i5
+```
+echo "export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:~/mirage/xembody/xembody/src/ros_ws/build/gazebo_env" >> ~/.bashrc
+```
+
+### Step 5: Compile workspace
+```
+cd ~/mirage/xembody/xembody/src/ros_ws
+colcon build
+. install/setup.bash
+source ~/.bashrc
+```
 When doing this installation, have conda fully deactivated
 Make sure ROS2 Humble Desktop Full is installed: https://docs.ros.org/en/humble/Installation/Ubuntu-Install-Debians.html