From ce6734e25824bb64a9d27044835667e63022d6fc Mon Sep 17 00:00:00 2001
From: wjxgeorge <12130815d@connect.polyu.hk>
Date: Mon, 25 Nov 2024 23:11:27 -0500
Subject: [PATCH] update Fetch and Robotic environments to fix reset issue
 introduced in v1.3.0

---
 gymnasium_robotics/__init__.py             | 10 +++++-----
 gymnasium_robotics/envs/fetch/fetch_env.py |  6 ++++++
 gymnasium_robotics/envs/robot_env.py       |  8 ++++++++
 3 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/gymnasium_robotics/__init__.py b/gymnasium_robotics/__init__.py
index feb748a6..c1481ddc 100644
--- a/gymnasium_robotics/__init__.py
+++ b/gymnasium_robotics/__init__.py
@@ -30,7 +30,7 @@ def _merge(a, b):
         )
 
         register(
-            id=f"FetchSlide{suffix}-v3",
+            id=f"FetchSlide{suffix}-v4",
             entry_point="gymnasium_robotics.envs.fetch.slide:MujocoFetchSlideEnv",
             kwargs=kwargs,
             max_episode_steps=50,
@@ -44,7 +44,7 @@ def _merge(a, b):
         )
 
         register(
-            id=f"FetchPickAndPlace{suffix}-v3",
+            id=f"FetchPickAndPlace{suffix}-v4",
             entry_point="gymnasium_robotics.envs.fetch.pick_and_place:MujocoFetchPickAndPlaceEnv",
             kwargs=kwargs,
             max_episode_steps=50,
@@ -58,7 +58,7 @@ def _merge(a, b):
         )
 
         register(
-            id=f"FetchReach{suffix}-v3",
+            id=f"FetchReach{suffix}-v4",
             entry_point="gymnasium_robotics.envs.fetch.reach:MujocoFetchReachEnv",
             kwargs=kwargs,
             max_episode_steps=50,
@@ -72,7 +72,7 @@ def _merge(a, b):
         )
 
         register(
-            id=f"FetchPush{suffix}-v3",
+            id=f"FetchPush{suffix}-v4",
             entry_point="gymnasium_robotics.envs.fetch.push:MujocoFetchPushEnv",
             kwargs=kwargs,
             max_episode_steps=50,
@@ -87,7 +87,7 @@ def _merge(a, b):
         )
 
         register(
-            id=f"HandReach{suffix}-v2",
+            id=f"HandReach{suffix}-v3",
             entry_point="gymnasium_robotics.envs.shadow_dexterous_hand.reach:MujocoHandReachEnv",
             kwargs=kwargs,
             max_episode_steps=50,
diff --git a/gymnasium_robotics/envs/fetch/fetch_env.py b/gymnasium_robotics/envs/fetch/fetch_env.py
index 1e90984c..5b59c2c6 100644
--- a/gymnasium_robotics/envs/fetch/fetch_env.py
+++ b/gymnasium_robotics/envs/fetch/fetch_env.py
@@ -376,6 +376,12 @@ def _reset_sim(self):
         # Reset buffers for joint states, actuators, warm-start, control buffers etc.
         self._mujoco.mj_resetData(self.model, self.data)
 
+        self.data.time = self.initial_time
+        self.data.qpos[:] = np.copy(self.initial_qpos)
+        self.data.qvel[:] = np.copy(self.initial_qvel)
+        if self.model.na != 0:
+            self.data.act[:] = None
+
         # Randomize start position of object.
         if self.has_object:
             object_xpos = self.initial_gripper_xpos[:2]
diff --git a/gymnasium_robotics/envs/robot_env.py b/gymnasium_robotics/envs/robot_env.py
index e482715d..9353873a 100644
--- a/gymnasium_robotics/envs/robot_env.py
+++ b/gymnasium_robotics/envs/robot_env.py
@@ -305,6 +305,14 @@ def _initialize_simulation(self):
     def _reset_sim(self):
         # Reset buffers for joint states, warm-start, control buffers etc.
         mujoco.mj_resetData(self.model, self.data)
+
+        self.data.time = self.initial_time
+        self.data.qpos[:] = np.copy(self.initial_qpos)
+        self.data.qvel[:] = np.copy(self.initial_qvel)
+        if self.model.na != 0:
+            self.data.act[:] = None
+
+        self._mujoco.mj_forward(self.model, self.data)
         return super()._reset_sim()
 
     def render(self):