UPGRADE: Adding franka-kitchen randomization to the RANDOM_ENTRY_POIN…

…T. This is what was called RANDOMDESK_ENTRY_POINT in the R3M codebase. This randomization is important to study visual generalizaiton. Otherwise proprioception alone is enough to solve the tasks

mj_envs/envs/multi_task/common/franka_kitchen_v1.py

@@ -197,12 +197,18 @@ def __init__(self, model_path, obsd_model_path=None, seed=None, **kwargs):
 
         super()._setup(**kwargs)
 
-    def reset(self, reset_qpos=None, reset_qvel=None):
+    def reset(self, reset_qpos=None, reset_qvel=None, reset_franka_wrt_kitchen=True):
         if reset_qpos is None:
             reset_qpos = self.init_qpos.copy()
             reset_qpos[self.robot_dofs] += (
                 0.05
                 * (self.np_random.uniform(size=len(self.robot_dofs)) - 0.5)
                 * (self.robot_ranges[:, 1] - self.robot_ranges[:, 0])
             )
-        return super().reset(reset_qpos=reset_qpos, reset_qvel=reset_qvel)
+
+        # reset franka wrt kitchen
+        if reset_franka_wrt_kitchen:
+            bid = self.sim.model.body_name2id("chef")
+            self.sim.model.body_pos[bid] = np.array([0, 0, 1.8]) + self.np_random.uniform(-0.1, 0.1, (3,))
+
+        return super().reset(reset_qpos=reset_qpos, reset_qvel=reset_qvel)