vikashplus / robohive

A unified framework for robot learning
https://sites.google.com/view/robohive
Apache License 2.0
515 stars 86 forks source link

set_env_state function not working #51

Closed ShahRutav closed 1 year ago

ShahRutav commented 1 year ago

On calling set_env_state I am getting an error:

Traceback (most recent call last):
  File "test.py", line 67, in <module>
    env.set_env_state(state)
  File "/home/rutavms/research/robohive/latest/mj_envs/mj_envs/envs/env_base.py", line 487, in set_env_state
    self.set_state(qp, qv, act)
AttributeError: 'FrankaAppliance' object has no attribute 'set_state'

I believe the set_env_state() function should be changed to:

     def set_env_state(self, state_dict):
       """
       Set full state of the environemnt
       Default implemention provided. Override if env has custom state
       """
          qp = state_dict['qpos']
          qv = state_dict['qvel']
          act = state_dict['act']
          self.sim.set_state(qpos=qp, qvel=qv, act=act)
          self.sim_obsd.set_state(qpos=qp, qvel=qv, act=act)
          if self.sim.model.nmocap>0:
              self.sim.model.mocap_pos[:] = state_dict['mocap_pos']
              self.sim.model.mocap_quat[:] = state_dict['mocap_quat']
              self.sim_obsd.model.mocap_pos[:] = state_dict['mocap_pos']
              self.sim_obsd.model.mocap_quat[:] = state_dict['mocap_quat']
          if self.sim.model.nsite>0:
              self.sim.model.site_pos[:] = state_dict['site_pos']
              self.sim.model.site_quat[:] = state_dict['site_quat']
              self.sim_obsd.model.site_pos[:] = state_dict['site_pos']
              self.sim_obsd.model.site_quat[:] = state_dict['site_quat']
          self.sim.model.body_pos[:] = state_dict['body_pos']
          self.sim.model.body_quat[:] = state_dict['body_quat']
          self.sim.forward()
          self.sim_obsd.model.body_pos[:] = state_dict['body_pos']
          self.sim_obsd.model.body_quat[:] = state_dict['body_quat']
          self.sim_obsd.forward()
vikashplus commented 1 year ago

fixed in 3318cc95f79c5df56f1bf346fa95bd754f85633f