move goal site to end effector on init

robohive/envs/multi_task/common/kitchen/franka_kitchen.xml

@@ -47,7 +47,7 @@
             <geom type="box" group="3" pos='0 0 .142' size="0.02 0.10 0.03" contype="0" conaffinity="0" rgba=".9 .7 .95 1" euler="0 0 -.785"/>
         </body>
 
-        <site name='target' type='box' size='.03 .07 .04' pos='-0.4 0.4 2.1' group='1' rgba='0 1 .4 0' euler="0 3.14 3.14"/>
+        <site name='target' type='box' size='.03 .07 .04' pos='0 0 0' group='1' rgba='0 1 .4 0' euler="0 3.14 3.14"/>
         <camera name='left_cam' pos='-1.4 -0.75 3' quat='0.78 0.49 -0.22 -0.32' />
         <camera name='right_cam' pos='1.4 -0.75 3' quat='0.76 0.5 0.21 0.35'/>
         <camera name='top_cam' pos='-.3 -.95 3.5' euler='0.785 0 0'/>

robohive/tutorials/ee_teleop.py

@@ -13,7 +13,7 @@
 """
 # TODO: (1) Enforce pos/rot/grip limits (b) move gripper to delta commands
 
-from robohive.utils.quat_math import euler2quat, mulQuat
+from robohive.utils.quat_math import euler2quat, mulQuat, mat2quat
 from robohive.utils.inverse_kinematics import IKResult, qpos_from_site_pose
 from robohive.logger.roboset_logger import RoboSet_Trace
 from robohive.logger.grouped_datasets import Trace as RoboHive_Trace
@@ -148,6 +148,21 @@ def poll_gamepad(input_device):
 
     return delta_pos * scale_factor, delta_euler * scale_factor, delta_gripper, done
 
+def move_goal_site_to_end_effector(teleop_site, goal_site, physics):
+    """
+    Get the location of the teleop site (the end effector),
+    and place the goal site exactly at this location.
+
+    teleop_site: A string specifying the name of the teleoperation site.
+    goal_site: A string specifying the name of the goal site.
+    physics: A `mujoco.Physics` instance.
+    """
+    ee_sid = physics.model.site_name2id(teleop_site)
+    goal_sid = physics.model.site_name2id(goal_site)
+    ee_xpos = physics.data.site_xpos[ee_sid]
+    ee_xquat = mat2quat(physics.data.site_xmat[ee_sid].reshape(3,3))
+    physics.model.site_pos[goal_sid] = ee_xpos
+    physics.model.site_quat[goal_sid] = ee_xquat
 
 @click.command(help=DESC)
 @click.option('-e', '--env_name', type=str, help='environment to load', default='rpFrankaRobotiqData-v0')
@@ -216,6 +231,9 @@ def main(env_name, env_args, reset_noise, action_noise, input_device, output, ho
         act = np.zeros(env.action_space.shape)
         gripper_state = 0
 
+        # Position the goal site exactly at the init location of the end effector
+        move_goal_site_to_end_effector(teleop_site, goal_site, env.sim)
+
         # start rolling out
         for i_step in range(horizon+1):