Set the bounded poses for the objects and target like panda-gym

gym_lowcostrobot/envs/base_env.py

@@ -129,3 +129,11 @@ def get_camera_images(self):
             img = self.renderer.render()
             dict_cams[cam_ids] = img
         return dict_cams
+
+    def set_object_range(self, obj_xy_range):
+        self.object_low = np.array([-obj_xy_range, -obj_xy_range, 0.01])
+        self.object_high = np.array([obj_xy_range, obj_xy_range, 0.01])
+
+    def set_target_range(self, target_xy_range):
+        self.target_low = np.array([-target_xy_range, -target_xy_range, 0.01])
+        self.target_high = np.array([target_xy_range, target_xy_range, 0.01])

gym_lowcostrobot/envs/lift_cube_env.py

@@ -9,7 +9,11 @@
 
 
 class LiftCubeEnv(BaseRobotEnv):
-    def __init__(self, image_state=None, action_mode="joint", render_mode=None):
+    def __init__(self,
+                 image_state=None, 
+                 action_mode="joint", 
+                 render_mode=None,
+                 obj_xy_range=0.15):
         super().__init__(
             xml_path="assets/scene_one_cube.xml",
             image_state=image_state,
@@ -30,9 +34,7 @@ def __init__(self, image_state=None, action_mode="joint", render_mode=None):
         )
 
         self.threshold_distance = 0.5
-
-        self.object_low = np.array([0.0, 0.0, 0.01])
-        self.object_high = np.array([0.2, 0.2, 0.01])
+        self.set_object_range(obj_xy_range)
 
     def reset(self, seed=None, options=None):
         # We need the following line to seed self.np_random

gym_lowcostrobot/envs/pick_place_cube_env.py

@@ -9,7 +9,11 @@
 
 
 class PickPlaceCubeEnv(BaseRobotEnv):
-    def __init__(self, image_state=None, action_mode="joint", render_mode=None):
+    def __init__(self, image_state=None, 
+                 action_mode="joint", 
+                 render_mode=None,
+                 target_xy_range=0.2,
+                 obj_xy_range=0.2):
         super().__init__(
             xml_path="assets/scene_one_cube.xml",
             image_state=image_state,
@@ -30,11 +34,8 @@ def __init__(self, image_state=None, action_mode="joint", render_mode=None):
         )
 
         self.threshold_distance = 0.26
-
-        self.target_low = np.array([0.0, 0.0, 0.01])
-        self.target_high = np.array([0.2, 0.2, 0.01])
-        self.object_low = np.array([0.0, 0.0, 0.01])
-        self.object_high = np.array([0.2, 0.2, 0.01])
+        self.set_object_range(obj_xy_range)
+        self.set_target_range(target_xy_range)
 
     def reset(self, seed=None, options=None):
         # We need the following line to seed self.np_random

gym_lowcostrobot/envs/push_cube_env.py

@@ -9,7 +9,11 @@
 
 
 class PushCubeEnv(BaseRobotEnv):
-    def __init__(self, image_state=None, action_mode="joint", render_mode=None):
+    def __init__(self, image_state=None, 
+                 action_mode="joint", 
+                 render_mode=None,
+                 target_xy_range=0.2,
+                 obj_xy_range=0.2):
         super().__init__(
             xml_path="assets/scene_one_cube.xml",
             image_state=image_state,
@@ -30,11 +34,8 @@ def __init__(self, image_state=None, action_mode="joint", render_mode=None):
         )
 
         self.threshold_distance = 0.26
-
-        self.target_low = np.array([0.0, 0.0, 0.01])
-        self.target_high = np.array([0.2, 0.2, 0.01])
-        self.object_low = np.array([0.0, 0.0, 0.01])
-        self.object_high = np.array([0.2, 0.2, 0.01])
+        self.set_object_range(obj_xy_range)
+        self.set_target_range(target_xy_range)
 
     def reset(self, seed=None, options=None):
         # We need the following line to seed self.np_random

gym_lowcostrobot/envs/reach_cube_env.py

@@ -9,7 +9,11 @@
 
 
 class ReachCubeEnv(BaseRobotEnv):
-    def __init__(self, image_state=None, action_mode="joint", render_mode=None):
+    def __init__(self,
+                 image_state=None, 
+                 action_mode="joint", 
+                 render_mode=None,
+                 obj_xy_range=0.15):
         super().__init__(
             xml_path="assets/scene_one_cube.xml",
             image_state=image_state,
@@ -31,9 +35,7 @@ def __init__(self, image_state=None, action_mode="joint", render_mode=None):
 
         # Initialize the robot and target positions
         self.threshold_distance = 0.01
-
-        self.object_low = np.array([0.0, 0.0, 0.01])
-        self.object_high = np.array([0.2, 0.2, 0.01])
+        self.set_object_range(obj_xy_range)
 
     def reset(self, seed=None, options=None):
         # We need the following line to seed self.np_random

gym_lowcostrobot/envs/stack_env.py

@@ -9,7 +9,11 @@
 
 
 class StackEnv(BaseRobotEnv):
-    def __init__(self, image_state=None, action_mode="joint", render_mode=None):
+    def __init__(self,
+                 image_state=None, 
+                 action_mode="joint", 
+                 render_mode=None,
+                 obj_xy_range=0.15):
         super().__init__(
             xml_path="assets/scene_two_cubes.xml",
             image_state=image_state,
@@ -30,9 +34,7 @@ def __init__(self, image_state=None, action_mode="joint", render_mode=None):
         )
 
         self.threshold_distance = 0.05
-
-        self.object_low = np.array([0.0, 0.0, 0.01])
-        self.object_high = np.array([0.2, 0.2, 0.01])
+        self.set_object_range(obj_xy_range)
 
     def reset(self, seed=None, options=None):
         # We need the following line to seed self.np_random