BUGFIX: Adding kwargs to all resets

robohive/envs/claws/reorient_v0.py

@@ -98,7 +98,7 @@ 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):
+ def reset(self, **kwargs):
  desired_pos = self.np_random.uniform(high=self.target_xyz_range['high'], low=self.target_xyz_range['low'])
  self.sim.model.site_pos[self.target_sid] = desired_pos
  self.sim_obsd.model.site_pos[self.target_sid] = desired_pos
@@ -108,5 +108,5 @@ def reset(self):
  self.sim.model.site_quat[self.target_sid] = euler2quat(desired_orien)
  self.sim_obsd.model.site_quat[self.target_sid] = euler2quat(desired_orien)
 
- obs = super().reset(self.init_qpos, self.init_qvel)
+ obs = super().reset(self.init_qpos, self.init_qvel, **kwargs)
  return obs

robohive/envs/fm/franka_ee_pose_v0.py

@@ -119,9 +119,9 @@ def get_target_pose(self):
  return self.np_random.uniform(low=self.sim.model.actuator_ctrlrange[:,0], high=self.sim.model.actuator_ctrlrange[:,1])
 
 
- def reset(self, reset_qpos=None, reset_qvel=None):
+ def reset(self, reset_qpos=None, reset_qvel=None, **kwargs):
  self.target_pose = self.get_target_pose()
- obs = super().reset(reset_qpos, reset_qvel)
+ obs = super().reset(reset_qpos, reset_qvel, **kwargs)
  return obs
 
 class FrankaRobotiqPose(FrankaEEPose):

robohive/envs/hands/baoding_v1.py

@@ -256,15 +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):
+ 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)
+ obs = super().reset(reset_qpos=reset_pose, reset_qvel=reset_vel, **kwargs)
  return obs
 
  def create_goal_trajectory(self, time_step=.1, time_period=6):
@@ -325,6 +325,6 @@ def create_goal_trajectory(self, time_step=.1, time_period=6):
 
 class BaodingRandomEnvV1(BaodingFixedEnvV1):
 
- def reset(self):
- obs = super().reset(time_period = self.np_random.uniform(high=5, low=7))
+ def reset(self, **kwargs):
+ obs = super().reset(time_period = self.np_random.uniform(high=5, low=7), **kwargs)
  return obs

robohive/envs/multi_task/common/franka_appliance_v1.py

@@ -54,7 +54,7 @@ def _setup(
  **kwargs,
  )
 
- def reset(self, reset_qpos=None, reset_qvel=None):
+ def reset(self, reset_qpos=None, reset_qvel=None, **kwargs):
  # randomize object bodies, if requested
  if self.obj_body_randomize:
  for body_name in self.obj_body_randomize:
@@ -78,4 +78,4 @@ def reset(self, reset_qpos=None, reset_qvel=None):
  * (self.robot_ranges[:, 1] - self.robot_ranges[:, 0])
  )
 
- return super().reset(reset_qpos=reset_qpos, reset_qvel=reset_qvel)
+ return super().reset(reset_qpos=reset_qpos, reset_qvel=reset_qvel, **kwargs)

robohive/envs/multi_task/common/franka_kitchen_v2.py

@@ -113,7 +113,7 @@ def _setup(
  )
 
 
- def reset(self, reset_qpos=None, reset_qvel=None):
+ def reset(self, reset_qpos=None, reset_qvel=None, **kwargs):
  if reset_qpos is None:
  reset_qpos = self.init_qpos.copy()
 
@@ -128,4 +128,4 @@ def reset(self, reset_qpos=None, reset_qvel=None):
  if self.robot_base_range:
  self.sim.model.body_pos[self.robot_base_bid] = self.robot_base_pos + self.np_random.uniform(**self.robot_base_range)
 
- return super().reset(reset_qpos=reset_qpos, reset_qvel=reset_qvel)
+ return super().reset(reset_qpos=reset_qpos, reset_qvel=reset_qvel, **kwargs)

robohive/envs/myo/myobase/baoding_v1.py

@@ -260,7 +260,7 @@ def evaluate_success(self, paths, logger=None, successful_steps=5):
  logger.log_kv('effort', effort)
  return success_percentage
 
- def reset(self, reset_pose=None, reset_vel=None, reset_goal=None, time_period=None):
+ def reset(self, reset_pose=None, reset_vel=None, reset_goal=None, time_period=None, **kwargs):
  # reset counters
  self.counter=0
  self.x_radius=self.np_random.uniform(low=self.goal_xrange[0], high=self.goal_xrange[1])
@@ -273,7 +273,7 @@ def reset(self, reset_pose=None, reset_vel=None, reset_goal=None, time_period=No
  self.goal = self.create_goal_trajectory(time_step=self.dt, 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)
+ obs = super().reset(reset_qpos=reset_pose, reset_qvel=reset_vel, **kwargs)
  return obs
 
  def create_goal_trajectory(self, time_step=.1, time_period=6):

robohive/envs/myo/myodm/myodm_v0.py

@@ -289,10 +289,10 @@ def playback(self):
  return idxs[0] < self.ref.horizon-1
 
 
- def reset(self):
+ def reset(self, **kwargs):
  # print("Reset")
  self.ref.reset()
- obs = super().reset(self.init_qpos, self.init_qvel)
+ obs = super().reset(self.init_qpos, self.init_qvel, **kwargs)
  # print(self.time, self.sim.data.qpos)
  return obs
 

robohive/envs/myo/myomimic/myomimic_v0.py

@@ -155,9 +155,9 @@ def playback(self):
  return idxs[0] < self.ref.horizon-1
 
 
- def reset(self):
+ def reset(self, **kwargs):
  # print("Reset")
  self.ref.reset()
- obs = super().reset(self.init_qpos, self.init_qvel)
+ obs = super().reset(self.init_qpos, self.init_qvel, **kwargs)
  # print(self.time, self.sim.data.qpos)
  return obs

robohive/envs/quadrupeds/orient_v0.py

@@ -171,7 +171,7 @@ def get_reward_dict(self, obs_dict):
  return rwd_dict
 
 
- def reset(self, reset_qpos=None, reset_qvel=None):
+ def reset(self, reset_qpos=None, reset_qvel=None, **kwargs):
 
  reset_qpos = self.init_qpos.copy() if reset_qpos is None else reset_qpos
  reset_qpos[6:] += np.pi/8*self.np_random.uniform(low=-1, high=1, size=self.sim.model.nq-6)
@@ -182,5 +182,5 @@ def reset(self, reset_qpos=None, reset_qvel=None):
  self.sim.model.site_pos[self.target_sid] = target_dist * np.array([np.cos(target_theta), np.sin(target_theta), 0])
  # Heading target is a bit farther away to avoid heading oscillations when quad is near xy_target
  self.sim.model.site_pos[self.heading_sid] = (target_dist+0.5) * np.array([np.cos(target_theta), np.sin(target_theta), 0])
- obs = super().reset(reset_qpos, reset_qvel)
+ obs = super().reset(reset_qpos, reset_qvel, **kwargs)
  return obs

robohive/envs/quadrupeds/stand_v0.py

@@ -170,7 +170,7 @@ def get_reward_dict(self, obs_dict):
  return rwd_dict
 
 
- def reset(self, reset_qpos=None, reset_qvel=None):
+ def reset(self, reset_qpos=None, reset_qvel=None, **kwargs):
 
  if reset_qpos is None:
  reset_qpos = self.init_qpos.copy()
@@ -189,5 +189,5 @@ def reset(self, reset_qpos=None, reset_qvel=None):
  else:
  raise TypeError(f"Unknown reset type: {self.reset_type}")
 
- obs = super().reset(reset_qpos, reset_qvel)
+ obs = super().reset(reset_qpos, reset_qvel, **kwargs)
  return obs

robohive/envs/quadrupeds/walk_v0.py

@@ -171,7 +171,7 @@ def get_reward_dict(self, obs_dict):
  return rwd_dict
 
 
- def reset(self, reset_qpos=None, reset_qvel=None):
+ def reset(self, reset_qpos=None, reset_qvel=None, **kwargs):
 
  reset_qpos = self.init_qpos.copy() if reset_qpos is None else reset_qpos
  reset_qpos[6:] += np.pi/8*self.np_random.uniform(low=-1, high=1, size=self.sim.model.nq-6)
@@ -182,5 +182,5 @@ def reset(self, reset_qpos=None, reset_qvel=None):
  self.sim.model.site_pos[self.target_sid] = target_dist * np.array([np.cos(target_theta), np.sin(target_theta), 0])
  # Heading target is a bit farther away to avoid heading oscillations when quad is near xy_target
  self.sim.model.site_pos[self.heading_sid] = (target_dist+0.5) * np.array([np.cos(target_theta), np.sin(target_theta), 0])
- obs = super().reset(reset_qpos, reset_qvel)
+ obs = super().reset(reset_qpos, reset_qvel, **kwargs)
  return obs

robohive/envs/tcdm/track.py

@@ -277,10 +277,10 @@ def playback(self):
  return idxs[0] < self.ref.horizon-1
 
 
- def reset(self):
+ def reset(self, **kwargs):
  # print("Reset")
  self.ref.reset()
- obs = super().reset(self.init_qpos, self.init_qvel)
+ obs = super().reset(self.init_qpos, self.init_qvel, **kwargs)
  # print(self.time, self.sim.data.qpos)
  return obs