BUGFIX: Fixing return type issues of envs with custom reset

- gym/gymnasium compatibility doesn't port over when custom reset is provided
- we are now recommending a call to super().reset(**kwargs) to get the appropritae return type ensuring compatibility

robohive/envs/env_base.py

@@ -495,8 +495,8 @@ def get_input_seed(self):
 
  def _reset(self, reset_qpos=None, reset_qvel=None, seed=None, **kwargs):
  """
- Reset the environment
- Default implemention provided. Override if env needs custom reset
+ Reset the environment (Default implemention provided).
+ Override if env needs custom reset. Carefully handle return type for gym/gymnasium compatibility
  """
  qpos = self.init_qpos.copy() if reset_qpos is None else reset_qpos
  qvel = self.init_qvel.copy() if reset_qvel is None else reset_qvel

robohive/envs/hands/baoding_v1.py

@@ -256,16 +256,15 @@ def get_reward_dict(self, obs_dict):
  rwd_dict['dense'] = np.sum([wt*rwd_dict[key] for key, wt in self.rwd_keys_wt.items()], axis=0)
  return rwd_dict
 
- def _reset(self, reset_pose=None, reset_vel=None, reset_goal=None, time_period=6, **kwargs):
+ def reset(self, reset_pose=None, reset_vel=None, reset_goal=None, time_period=6, **kwargs):
  # reset counters
  self.counter=0
 
  # reset goal
  self.goal = self.create_goal_trajectory(time_period=time_period) if reset_goal is None else reset_goal.copy()
 
  # reset scene
- obs = super().reset(reset_qpos=reset_pose, reset_qvel=reset_vel, **kwargs)
- return obs
+ return super().reset(reset_qpos=reset_pose, reset_qvel=reset_vel, **kwargs)
 
  def create_goal_trajectory(self, time_step=.1, time_period=6):
  len_of_goals = 1000 # assumes that its greator than env horizon
@@ -326,5 +325,4 @@ def create_goal_trajectory(self, time_step=.1, time_period=6):
 class BaodingRandomEnvV1(BaodingFixedEnvV1):
 
  def reset(self, **kwargs):
- obs = super().reset(time_period = self.np_random.uniform(high=5, low=7), **kwargs)
- return obs
+ return super().reset(time_period = self.np_random.uniform(high=5, low=7), **kwargs)

robohive/envs/hands/door_v1.py

@@ -96,17 +96,13 @@ def get_reward_dict(self, obs_dict):
  return rwd_dict
 
 
- def _reset(self, reset_qpos=None, reset_qvel=None, **kwargs):
+ def reset(self, **kwargs):
  self.sim.reset()
- qp = self.init_qpos.copy() if reset_qpos is None else reset_qpos
- qv = self.init_qvel.copy() if reset_qvel is None else reset_qvel
- self.robot.reset(reset_pos=qp, reset_vel=qv, **kwargs)
-
  self.sim.model.body_pos[self.door_bid,0] = self.np_random.uniform(low=-0.3, high=-0.2)
  self.sim.model.body_pos[self.door_bid, 1] = self.np_random.uniform(low=0.25, high=0.35)
  self.sim.model.body_pos[self.door_bid,2] = self.np_random.uniform(low=0.252, high=0.35)
  self.sim.forward()
- return self.get_obs()
+ return super().reset(**kwargs)
 
 
  def get_env_state(self):

robohive/envs/hands/hammer_v1.py

@@ -120,15 +120,12 @@ def get_obs_dict(self, sim):
  return obs_dict
 
 
- def _reset(self, reset_qpos=None, reset_qvel=None, **kwargs):
+ def reset(self, **kwargs):
  self.sim.reset()
- qp = self.init_qpos.copy() if reset_qpos is None else reset_qpos
- qv = self.init_qvel.copy() if reset_qvel is None else reset_qvel
- self.robot.reset(reset_pos=qp, reset_vel=qv, **kwargs)
-
  self.sim.model.body_pos[self.target_bid,2] = self.np_random.uniform(low=0.1, high=0.25)
  self.sim.forward()
- return self.get_obs()
+ return super().reset(**kwargs)
+
 
  def get_env_state(self):
  """

robohive/envs/hands/pen_v1.py

@@ -108,19 +108,14 @@ def get_reward_dict(self, obs_dict):
  return rwd_dict
 
 
- def _reset(self, reset_qpos=None, reset_qvel=None, **kwargs):
+ def reset(self, **kwargs):
  self.sim.reset()
- qp = self.init_qpos.copy() if reset_qpos is None else reset_qpos
- qv = self.init_qvel.copy() if reset_qvel is None else reset_qvel
- self.robot.reset(reset_pos=qp, reset_vel=qv, **kwargs)
-
  desired_orien = np.zeros(3)
  desired_orien[0] = self.np_random.uniform(low=-1, high=1)
  desired_orien[1] = self.np_random.uniform(low=-1, high=1)
  self.sim.model.body_quat[self.target_obj_bid] = euler2quat(desired_orien)
  self.sim.forward()
-
- return self.get_obs()
+ return super().reset(**kwargs)
 
 
  def get_env_state(self):

robohive/envs/hands/relocate_v1.py

@@ -136,20 +136,16 @@ def get_obs_dict(self, sim):
  return obs_dict
 
 
- def _reset(self, reset_qpos=None, reset_qvel=None, **kwargs):
+ def reset(self, **kwargs):
  self.sim.reset()
- qp = self.init_qpos.copy() if reset_qpos is None else reset_qpos
- qv = self.init_qvel.copy() if reset_qvel is None else reset_qvel
- self.robot.reset(reset_pos=qp, reset_vel=qv, **kwargs)
-
-
  self.sim.model.body_pos[self.obj_bid,0] = self.np_random.uniform(low=-0.15, high=0.15)
  self.sim.model.body_pos[self.obj_bid,1] = self.np_random.uniform(low=-0.15, high=0.3)
  self.sim.model.site_pos[self.target_obj_sid, 0] = self.np_random.uniform(low=-0.2, high=0.2)
  self.sim.model.site_pos[self.target_obj_sid,1] = self.np_random.uniform(low=-0.2, high=0.2)
  self.sim.model.site_pos[self.target_obj_sid,2] = self.np_random.uniform(low=0.15, high=0.35)
  self.sim.forward()
- return self.get_obs()
+ return super().reset(**kwargs)
+
 
  def get_env_state(self):
  """

robohive/tests/test_envs.py

@@ -8,6 +8,7 @@
 import unittest
 
 from robohive.utils import gym
+from robohive.utils.implement_for import implement_for
 import numpy as np
 import pickle
 import copy
@@ -66,7 +67,8 @@ def check_env(self, environment_id, input_seed):
  rwd_dict1 = env1.get_reward_dict(obs_dict1)
  assert len(rwd_dict1) > 0
  # reset env
- env1.reset()
+ reset_data = env1.reset()
+ self.check_reset(reset_data)
 
  # serialize / deserialize env ------------
  env2w = pickle.loads(pickle.dumps(env1w))
@@ -102,6 +104,20 @@ def check_env(self, environment_id, input_seed):
  del(env1)
  del(env2)
 
+
+ @implement_for("gym", None, "0.26")
+ def check_reset(self, reset_data):
+ assert isinstance(reset_data, np.ndarray), "Reset should return the observation vector"
+
+ @implement_for("gym", "0.26", None)
+ def check_reset(self, reset_data):
+ assert isinstance(reset_data, tuple) and len(reset_data) == 2, "Reset should return a tuple of length 2"
+ assert isinstance(reset_data[1], dict), "second element returned should be a dict"
+ @implement_for("gymnasium")
+ def check_reset(self, reset_data):
+ assert isinstance(reset_data, tuple) and len(reset_data) == 2, "Reset should return a tuple of length 2"
+ assert isinstance(reset_data[1], dict), "second element returned should be a dict"
+
  def check_old_envs(self, module_name, env_names, lite=False, seed=1234):
  print("\nTesting module:: ", module_name)
  for env_name in env_names: