Pickling works by calling mj_forward

Farama-Foundation · Jul 5, 2023 · 4e247bc · 4e247bc
1 parent cfe594e
commit 4e247bc
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 3 deletions.
diff --git a/metaworld/__init__.py b/metaworld/__init__.py
@@ -93,6 +93,7 @@ def _make_tasks(classes, args_kwargs, kwargs_override, seed=None):
         del kwargs["task_id"]
         env._set_task_inner(**kwargs)
         for _ in range(_N_GOALS):
+            print('reset')
             env.reset()
             rand_vecs.append(env._last_rand_vec)
         unique_task_rand_vecs = np.unique(np.array(rand_vecs), axis=0)

diff --git a/metaworld/envs/mujoco/sawyer_xyz/sawyer_xyz_env.py b/metaworld/envs/mujoco/sawyer_xyz/sawyer_xyz_env.py
@@ -33,7 +33,6 @@ def __init__(self, model_name, frame_skip=5):
         mjenv_gym.__init__(self, model_name, frame_skip=frame_skip, observation_space=self.observation_space)
         self.reset_mocap_welds()
         self.frame_skip = frame_skip
-        print(self.data)
 
     def get_endeff_pos(self):
         return self.data.body('hand').xpos
@@ -60,12 +59,14 @@ def set_env_state(self, state):
         self.set_state(mocap_pos, mocap_quat)
 
     def __getstate__(self):
+        print('__getstate__')
         state = self.__dict__.copy()
         # del state['model']
         # del state['data']
         return {'state': state, 'mjb': self.model_name, 'mocap':self.get_env_state()}
 
     def __setstate__(self, state):
+        print('__setstate__')
         self.__dict__ = state['state']
         mjenv_gym.__init__(self, state['mjb'], frame_skip=self.frame_skip, observation_space=self.observation_space)
         self.set_env_state(state['mocap'])
@@ -132,6 +133,9 @@ def __init__(
         self._partially_observable = True
 
         super().__init__(model_name, frame_skip=frame_skip)
+
+        mujoco.mj_forward(self.model, self.data)  # *** DO NOT REMOVE: EZPICKLE WON'T WORK *** #
+
         self._did_see_sim_exception = False
         self.init_left_pad = self.get_body_com("leftpad")
         self.init_right_pad = self.get_body_com("rightpad")
@@ -148,7 +152,6 @@ def __init__(
 
         self._set_task_called = False
 
-
         self.hand_init_pos = None  # OVERRIDE ME
         self._target_pos = None  # OVERRIDE ME
         self._random_reset_space = None  # OVERRIDE ME
@@ -162,6 +165,9 @@ def __init__(
 
         self._prev_obs = self._get_curr_obs_combined_no_goal()
         self.current_task = 0
+
+
+
         EzPickle.__init__(
             self,
             model_name,
@@ -173,6 +179,7 @@ def __init__(
             action_scale,
             action_rot_scale,
         )
+
     @staticmethod
     def _set_task_inner():
         # Doesn't absorb "extra" kwargs, to ensure nothing's missed.

diff --git a/metaworld/envs/mujoco/sawyer_xyz/v2/sawyer_reach_v2.py b/metaworld/envs/mujoco/sawyer_xyz/v2/sawyer_reach_v2.py
@@ -62,7 +62,6 @@ def __init__(self, tasks=None):
         self.goal_space = Box(np.array(goal_low), np.array(goal_high))
 
 
-
     @property
     def model_name(self):
         return full_v2_path_for("sawyer_xyz/sawyer_reach_v2.xml")