Update ENVmodel.py

ENVmodel.py

@@ -103,11 +103,12 @@ class SO101PickPlaceEnv(gym.Env):
 
     metadata = {"render_modes": ["human", "rgb_arrхay"], "render_fps": CONTROL_FREQ}
 
-    def __init__(self, render_mode="rgb_array", render_images=True):
+    def __init__(self, render_mode="rgb_array", render_images=True, max_episode_steps):
         super().__init__()
         self.render_mode = render_mode
         self.render_images = render_images
-
+        self.max_episode_steps = max_episode_steps
+       
         # Load MuJoCo model
         self.model = mujoco.MjModel.from_xml_path(SCENE_XML)
         self.model.opt.timestep = PHYSICS_TIMESTEP
@@ -261,7 +262,7 @@ class SO101PickPlaceEnv(gym.Env):
 
         reward = self._compute_reward(info)
         terminated = info["success"]
-        truncated = self._step_count >= MAX_EPISODE_STEPS
+        truncated = self._step_count >= self.max_episode_steps
 
         return obs, reward, terminated, truncated, info