Fix reproducibility of FetchPickAndPlace-v2

Farama-Foundation · Feb 10, 2024 · 54043fc · 54043fc
1 parent 50c8019
commit 54043fc
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 0 deletions.
diff --git a/gymnasium_robotics/envs/fetch/fetch_env.py b/gymnasium_robotics/envs/fetch/fetch_env.py
@@ -376,6 +376,10 @@ def _reset_sim(self):
         self.data.time = self.initial_time
         self.data.qpos[:] = np.copy(self.initial_qpos)
         self.data.qvel[:] = np.copy(self.initial_qvel)
+        self.data.qacc_warmstart[:] = np.copy(self.initial_qacc_warmstart)
+        self.data.ctrl[:] = np.copy(self.initial_ctrl)
+        self.data.mocap_pos[:] = np.copy(self.initial_mocap_pos)
+        self.data.mocap_quat[:] = np.copy(self.initial_mocap_quat)
         if self.model.na != 0:
             self.data.act[:] = None
 

diff --git a/gymnasium_robotics/envs/robot_env.py b/gymnasium_robotics/envs/robot_env.py
@@ -297,6 +297,10 @@ def _initialize_simulation(self):
         self.initial_time = self.data.time
         self.initial_qpos = np.copy(self.data.qpos)
         self.initial_qvel = np.copy(self.data.qvel)
+        self.initial_ctrl = np.copy(self.data.ctrl)
+        self.initial_qacc_warmstart = np.copy(self.data.qacc_warmstart)
+        self.initial_mocap_pos = np.copy(self.data.mocap_pos)
+        self.initial_mocap_quat = np.copy(self.data.mocap_quat)
 
     def _reset_sim(self):
         self.data.time = self.initial_time