BUGFIX: Extra index was added in np.choice option

vikashplus · Mar 25, 2024 · dc97c43 · dc97c43
1 parent 80ad5bf
commit dc97c43
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 4 deletions.
diff --git a/robohive/__init__.py b/robohive/__init__.py
@@ -12,13 +12,15 @@
 
 # Check for RoboHive initialization
 from robohive.utils.import_utils import simhive_isavailable
+
 simhive_isavailable(robohive_version = __version__)
 
 
 # Register RoboHive Envs
 from robohive.utils import gym
 from robohive.utils.implement_for import implement_for
 
+
 #TODO: check versions
 @implement_for("gym", None, "0.24")
 def gym_registry_specs():
@@ -32,8 +34,7 @@ def gym_registry_specs():
 def gym_registry_specs():
  return gym.envs.registry
 
-_current_gym_envs = gym_registry_specs().keys()
-_current_gym_envs = set(_current_gym_envs)
+_current_gym_envs = set(gym_registry_specs().keys())
 robohive_env_suite = set()
 
 # Register Arms Suite

diff --git a/robohive/envs/arms/pick_place_v0.py b/robohive/envs/arms/pick_place_v0.py
@@ -6,12 +6,14 @@
 ================================================= """
 
 import collections
-from robohive.utils import gym
+
 import numpy as np
 
 from robohive.envs import env_base
+from robohive.utils import gym
 from robohive.utils.quat_math import euler2quat
 
+
 class PickPlaceV0(env_base.MujocoEnv):
 
  DEFAULT_OBS_KEYS = [
@@ -130,7 +132,7 @@ def reset(self, **kwargs):
 
  for gid in range(self.sim.model.body_geomnum[bid]):
  gid+=self.sim.model.body_geomadr[bid]
- self.sim.model.geom_type[gid]=self.np_random.choice([2,3,4,5,6,7]) # random shape
+ self.sim.model.geom_type[gid]=self.np_random.choice([2,3,4,5,6]) # random shape
  self.sim.model.geom_size[gid]=self.np_random.uniform(low=self.geom_sizes['low'], high=self.geom_sizes['high']) # random size
  self.sim.model.geom_pos[gid]=self.np_random.uniform(low=-1*self.sim.model.geom_size[gid], high=self.sim.model.geom_size[gid]) # random pos
  self.sim.model.geom_quat[gid]=euler2quat(self.np_random.uniform(low=(-np.pi/2, -np.pi/2, -np.pi/2), high=(np.pi/2, np.pi/2, np.pi/2)) ) # random quat