Fix/reset bug

deformable_gym/envs/base_env.py

@@ -67,11 +67,13 @@ def _create_robot(self):
 
     def _load_objects(self):
         """Load objects to PyBullet simulation."""
-        self.plane = pb.loadURDF("plane.urdf", (0, 0, 0), useFixedBase=1)
+        self.plane = pb.loadURDF("plane.urdf",
+                                 (0, 0, 0),
+                                 useFixedBase=1)
 
     def _hard_reset(self):
-        """Hard reset the PyBullet simulation and reload all objects. May be necessary, e.g., if soft-bodies in the
-        environment explode.
+        """Hard reset the PyBullet simulation and reload all objects. May be
+        necessary, e.g., if soft-bodies in the environment explode.
         """
         if self.verbose:
             print("Performing hard reset!")
@@ -87,6 +89,9 @@ def reset(self, seed=None, options=None) -> npt.ArrayLike:
         """
         super().reset(seed=seed)
 
+        if options is not None and "hard_reset" in options:
+            self._hard_reset()
+
         assert isinstance(self.robot, BulletRobot)
 
         self.robot.reset()
@@ -131,7 +136,11 @@ def _get_info(self):
         """
         return {}
 
-    def _is_terminated(self, observation: npt.ArrayLike, action: npt.ArrayLike, next_observation: npt.ArrayLike) -> bool:
+    def _is_terminated(self,
+                       observation: npt.ArrayLike,
+                       action: npt.ArrayLike,
+                       next_observation: npt.ArrayLike
+                       ) -> bool:
         """Checks whether the current episode is terminated.
 
         :param observation: observation before action was taken
@@ -141,7 +150,10 @@ def _is_terminated(self, observation: npt.ArrayLike, action: npt.ArrayLike, next
         """
         return self.step_counter >= self.horizon
 
-    def _is_truncated(self, state: npt.ArrayLike, action: npt.ArrayLike, next_state: npt.ArrayLike) -> bool:
+    def _is_truncated(self,
+                      state: npt.ArrayLike,
+                      action: npt.ArrayLike,
+                      next_state: npt.ArrayLike) -> bool:
         """Checks whether the current episode is truncated.
 
         :param state: State

deformable_gym/envs/floating_mia_grasp_env.py

@@ -34,10 +34,10 @@ class FloatingMiaGraspEnv(GraspDeformableMixin, BaseBulletEnv):
         object should be sampled from the training or the test set.
     """
 
-    _STANDARD_INITIAL_POSE = np.r_[0.03, -0.005, 1.0,
+    STANDARD_INITIAL_POSE = np.r_[0.03, -0.005, 1.0,
                                    pb.getQuaternionFromEuler([-np.pi/8, np.pi, 0])]
 
-    _HARD_INITIAL_POSE = np.r_[0.03, -0.025, 1.0,
+    HARD_INITIAL_POSE = np.r_[0.03, -0.025, 1.0,
                                pb.getQuaternionFromEuler([-np.pi/8, np.pi, 0])]
 
     _FINGERS_OPEN = {"j_index_fle": 0.0,
@@ -58,13 +58,15 @@ class FloatingMiaGraspEnv(GraspDeformableMixin, BaseBulletEnv):
                        "j_ring_fle": 0.71,
                        "j_thumb_fle": 0.3}
 
+    _MAX_POS_OFFSET = .0005
+    _MAX_ORN_OFFSET = .000005
+
     def __init__(
             self,
             object_name: str = "insole",
             compute_reward: bool = True,
             object_scale: float = 1.0,
             observable_object_pos: bool = False,
-            observable_time_step: bool = False,
             difficulty_mode: str = "hard",
             **kwargs):
 
@@ -75,11 +77,10 @@ def __init__(
         self.compute_reward = compute_reward
         self.object_scale = object_scale
         self._observable_object_pos = observable_object_pos
-        self._observable_time_step = observable_time_step
 
         super().__init__(soft=True, **kwargs)
 
-        self.hand_world_pose = self._STANDARD_INITIAL_POSE
+        self.hand_world_pose = self.STANDARD_INITIAL_POSE
         self.robot = self._create_robot()
 
         limits = pbh.get_limit_array(self.robot.motors.values())
@@ -91,21 +92,19 @@ def __init__(
             np.array([-2, -2, 0]),
             -np.ones(4),
             limits[0][self.actuated_finger_ids],
-            -np.ones(6)*10])
+            -np.full(6, 10)])
 
         upper_observations = np.concatenate([
             np.array([2, 2, 2]),
             np.ones(4),
             limits[1][self.actuated_finger_ids],
-            np.ones(6)*10])
+            np.full(6, 10.)])
 
         if self._observable_object_pos:
-            lower_observations = np.append(lower_observations, -np.ones(3)*2)
-            upper_observations = np.append(upper_observations, np.ones(3)*2)
-
-        if self._observable_time_step:
-            lower_observations = np.append(lower_observations, 0)
-            upper_observations = np.append(upper_observations, self.horizon)
+            lower_observations = np.append(lower_observations,
+                                           -np.full(3, 2.))
+            upper_observations = np.append(upper_observations,
+                                           np.full(3, 2.))
 
         self.observation_space = spaces.Box(
             low=lower_observations,
@@ -114,12 +113,12 @@ def __init__(
         )
 
         # build the action space
-        lower = [-np.ones(3) * .0005,  # max negative base pos offset
-                 -np.ones(4) * .000005,  # max negative base orn offset
+        lower = [-np.full(3, self._MAX_POS_OFFSET),  # max negative base pos offset
+                 -np.full(4, self._MAX_ORN_OFFSET),  # max negative base orn offset
                  limits[0][self.actuated_finger_ids]]  # negative joint limits
 
-        upper = [np.ones(3) * .0005,  # max positive base pos offset
-                 np.ones(4) * .000005,  # max positive base orn offset
+        upper = [np.full(3, self._MAX_POS_OFFSET),  # max positive base pos offset
+                 np.full(4, self._MAX_ORN_OFFSET),  # max positive base orn offset
                  limits[1][self.actuated_finger_ids]]  # positive joint limits
 
         self.action_space = spaces.Box(low=np.concatenate(lower),
@@ -162,22 +161,22 @@ def set_difficulty_mode(self, mode: str):
 
         if mode == "hard":
             self.robot.set_initial_joint_positions(self._FINGERS_OPEN)
-            self.hand_world_pose = self._HARD_INITIAL_POSE
+            self.hand_world_pose = self.HARD_INITIAL_POSE
         elif mode == "medium":
             self.robot.set_initial_joint_positions(self._FINGERS_HALFWAY_CLOSED)
-            self.hand_world_pose = self._STANDARD_INITIAL_POSE
+            self.hand_world_pose = self.STANDARD_INITIAL_POSE
         elif mode == "easy":
             self.robot.set_initial_joint_positions(self._FINGERS_CLOSED)
-            self.hand_world_pose = self._STANDARD_INITIAL_POSE
+            self.hand_world_pose = self.STANDARD_INITIAL_POSE
         else:
             raise ValueError(f"Received unknown difficulty mode {mode}!")
 
-    def reset(self, seed=None, options=None, hard_reset=False, pos=None):
+    def reset(self, seed=None, options=None):
 
-        if pos is None:
-            self.robot.reset_base(self.hand_world_pose)
+        if options is not None and "initial_pose" in options:
+            self.robot.reset_base(options["initial_pose"])
         else:
-            self.robot.reset_base(pos)
+            self.robot.reset_base(self.hand_world_pose)
 
         self.object_to_grasp.reset()
         self.robot.activate_motors()
@@ -199,9 +198,6 @@ def _get_observation(self):
             obj_pos = self.object_to_grasp.get_pose()[:3]
             state = np.append(state, obj_pos)
 
-        if self._observable_time_step:
-            state = np.append(state, self.step_counter)
-
         return state
 
     def calculate_reward(self, state, action, next_state, terminated):

deformable_gym/envs/floating_shadow_grasp_env.py

@@ -33,7 +33,6 @@ def __init__(self, object_name="insole",
                  horizon=100, train=True,
                  compute_reward=True, object_scale=1.0,
                  observable_object_pos: bool = False,
-                 observable_time_step: bool = False,
                  **kwargs):
         self.insole = None
         self.train = train
@@ -43,7 +42,6 @@ def __init__(self, object_name="insole",
         self.compute_reward = compute_reward
         self.object_scale = object_scale
         self._observable_object_pos = observable_object_pos
-        self._observable_time_step = observable_time_step
 
         super().__init__(horizon=horizon, soft=True, **kwargs)
 
@@ -64,10 +62,6 @@ def __init__(self, object_name="insole",
             lower_observations = np.append(lower_observations, -np.ones(3)*2)
             upper_observations = np.append(upper_observations, np.ones(3)*2)
 
-        if self._observable_time_step:
-            lower_observations = np.append(lower_observations, 0)
-            upper_observations = np.append(upper_observations, self.horizon)
-
         self.observation_space = spaces.Box(
             low=lower_observations, high=upper_observations, dtype=np.float64)
 
@@ -104,7 +98,7 @@ def _load_objects(self):
         super()._load_objects()
         self.object_to_grasp, self.object_position, self.object_orientation = ObjectFactory().create(self.object_name)
 
-    def reset(self, seed=None, options=None, hard_reset=False):
+    def reset(self, seed=None, options=None):
         self.object_to_grasp.reset()
         self.robot.activate_motors()
 

deformable_gym/envs/grasp_env.py

@@ -9,12 +9,10 @@ class GraspEnv(BaseBulletEnv, ABC):
 
     def __init__(self,
                  object_name,
-                 observable_object_pos: bool = False,
-                 observable_time_step: bool = False):
+                 observable_object_pos: bool = False):
 
         self.object_name = object_name
         self._observable_object_pos = observable_object_pos
-        self._observable_time_step = observable_time_step
 
         super().__init__(soft=True)
 
@@ -28,13 +26,16 @@ def _load_objects(self):
         super()._load_objects()
         self.object_to_grasp, self.object_position, self.object_orientation = ObjectFactory().create(self.object_name)
 
-    def reset(self, hard_reset=False, pos=None):
+    def reset(self, seed=None, options=None):
+
+        if options is not None and "pos" in options:
+            pos = options["pos"]
 
         self.robot.reset(pos)
         self.object_to_grasp.reset()
         self.robot.activate_motors()
 
-        return super().reset()
+        return super().reset(seed, options)
 
     def _get_observation(self):
         joint_pos = self.robot.get_joint_positions(self.robot.actuated_real_joints)
@@ -46,9 +47,6 @@ def _get_observation(self):
             obj_pos = self.object_to_grasp.get_pose()[:3]
             state = np.append(state, obj_pos)
 
-        if self._observable_time_step:
-            state = np.append(state, self.step_counter)
-
         return state
 
     def calculate_reward(self, state, action, next_state, done):

deformable_gym/envs/ur10_shadow_grasp_env.py

@@ -97,14 +97,14 @@ def _load_objects(self):
         self.object_to_grasp, self.object_position, self.object_orientation = \
             ObjectFactory().create(self.object_name, object2world=self.object2world, scale=self.object_scale)
 
-    def reset(self, hard_reset=False):
+    def reset(self, seed=None, options=None):
 
         pos = None
 
         self.object_to_grasp.reset(pos)
         self.robot.activate_motors()
 
-        return super().reset()
+        return super().reset(seed, options)
 
     def is_done(self, state, action, next_state):
 

deformable_gym/envs/ur5_mia_grasp_env.py

@@ -92,11 +92,17 @@ class UR5MiaGraspEnv(GraspDeformableMixin, BaseBulletEnv):
                                      p=np.array([-0.7, 0.1, 1.8]))
 
     def __init__(
-            self, gui: bool = True, real_time: bool = False,
-            object_name: str = "insole", verbose: bool = False,
-            horizon: int = 100, train: bool = True,
-            thumb_adducted: bool = True, compute_reward: bool = True,
-            object_scale: float = 1.0, verbose_dt: float = 10.0,
+            self,
+            gui: bool = True,
+            real_time: bool = False,
+            object_name: str = "insole",
+            verbose: bool = False,
+            horizon: int = 100,
+            train: bool = True,
+            thumb_adducted: bool = True,
+            compute_reward: bool = True,
+            object_scale: float = 1.0,
+            verbose_dt: float = 10.0,
             pybullet_options: str = ""):
         self.train = train
         self.velocity_commands = False
@@ -166,7 +172,7 @@ def _load_objects(self):
         self.object_to_grasp, self.object_position, self.object_orientation = \
             ObjectFactory().create(self.object_name, object2world=self.object2world, scale=self.object_scale)
 
-    def reset(self, seed=None, options=None, hard_reset=False):
+    def reset(self, seed=None, options=None):
 
         pos = None
 

deformable_gym/helpers/pybullet_helper.py

@@ -20,8 +20,16 @@ class JointType(Enum):
 
 class Joint:
 
-    def __init__(self, name, joint_idx, joint_type, pos_idx, vel_idx, low, high,
-                 max_force, max_vel, body_id):
+    def __init__(self, name: str,
+                 joint_idx: int,
+                 joint_type: JointType,
+                 pos_idx: int,
+                 vel_idx: int,
+                 low: float,
+                 high: float,
+                 max_force: float,
+                 max_vel: float,
+                 body_id: int) -> None:
         self.name = name
         self.joint_idx = joint_idx
         self.joint_type = joint_type
@@ -40,29 +48,33 @@ def __init__(self, name, joint_idx, joint_type, pos_idx, vel_idx, low, high,
     def __repr__(self):
         return f"Joint({', '.join([f'{k}={v}' for k,v in vars(self).items()])})"
 
-    def reset(self):
+    def reset(self) -> None:
         """Resets the joint to its initial position."""
         self.set_position(self.init_pos)
 
-    def set_position(self, position):
+    def set_position(self, position: float) -> None:
         """Sets the joint to target position."""
-        pb.resetJointState(self.body_id, self.joint_idx,
-                           targetValue=position, targetVelocity=0.0)
+        pb.resetJointState(self.body_id,
+                           self.joint_idx,
+                           targetValue=position,
+                           targetVelocity=0.0)
 
-    def set_limits(self, low, high):
+    def set_limits(self, low: float, high: float) -> None:
         """Sets the joint limits."""
-        pb.changeDynamics(self.body_id, self.joint_idx,
-                          jointLowerLimit=low, jointUpperLimit=high)
+        pb.changeDynamics(self.body_id,
+                          self.joint_idx,
+                          jointLowerLimit=low,
+                          jointUpperLimit=high)
 
-    def get_position(self):
+    def get_position(self) -> float:
         """Gets the current position of the joint."""
         return pb.getJointState(self.body_id, self.joint_idx)[0]
 
-    def get_velocity(self):
+    def get_velocity(self) -> float:
         """Gets the current velocity of the joint."""
         return pb.getJointState(self.body_id, self.joint_idx)[1]
 
-    def set_target_position(self, position):
+    def set_target_position(self, position: float) -> None:
         """Sets the target position of the joint."""
 
         if self.max_vel + 0.1 < self.get_velocity():
@@ -82,7 +94,7 @@ def set_target_position(self, position):
             if self.verbose:
                 print(f"Warning: Trying to control deactivated motor {self.name}.")
 
-    def set_target_velocity(self, velocity):
+    def set_target_velocity(self, velocity: float) -> None:
         """Sets the target position of the joint."""
         if self.activated:
             pb.setJointMotorControl2(self.body_id, self.joint_idx,