Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Added get_pixel support in mujoco_env and proprioception in all envs #2

Open
wants to merge 5 commits into
base: old-master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from 4 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 13 additions & 0 deletions mj_envs/hand_manipulation_suite/door_v0.py
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,19 @@ def _step(self, a):

return ob, reward, False, {}

def get_proprioception(self, use_tactile):
# return self._get_obs()
robot_jnt = self.data.qpos.ravel()[:-2]
robot_vel = self.data.qvel.ravel()[:-2]
palm_pos = self.data.site_xpos[self.grasp_sid].ravel()
sensordata = []
if use_tactile:
sensordata = self.data.sensordata.ravel().copy()[:41]
sensordata = np.clip(sensordata, -5.0, 5.0)

res = np.concatenate([robot_jnt, robot_vel, palm_pos, sensordata])
return res

def _get_obs(self):
# qpos for hand
# xpos for obj
Expand Down
13 changes: 13 additions & 0 deletions mj_envs/hand_manipulation_suite/hammer_v0.py
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,19 @@ def _get_obs(self):
nail_impact = np.clip(self.sim.data.sensordata[self.sim.model.sensor_name2id('S_nail')], -1.0, 1.0)
return np.concatenate([qp[:-6], qv[-6:], palm_pos, obj_pos, obj_rot, target_pos, np.array([nail_impact])])

def get_proprioception(self, use_tactile):
# return self._get_obs()
robot_jnt = self.data.qpos.ravel()[:-6]
robot_vel = self.data.qvel.ravel()[:-6]
palm_pos = self.data.site_xpos[self.S_grasp_sid].ravel()
sensordata = []
if use_tactile:
sensordata = self.data.sensordata.ravel().copy()[:41]
sensordata = np.clip(sensordata, -5.0, 5.0)

res = np.concatenate([robot_jnt, robot_vel, palm_pos, sensordata])
return res

def reset_model(self):
self.sim.reset()
target_bid = self.model.body_name2id('nail_board')
Expand Down
12 changes: 12 additions & 0 deletions mj_envs/hand_manipulation_suite/pen_v0.py
Original file line number Diff line number Diff line change
Expand Up @@ -85,6 +85,18 @@ def _get_obs(self):
return np.concatenate([qp[:-6], obj_pos, obj_vel, obj_orien, desired_orien,
obj_pos-desired_pos, obj_orien-desired_orien])

def get_proprioception(self, use_tactile):
# return self._get_obs()
robot_jnt = self.data.qpos.ravel()[:-6]
robot_vel = self.data.qvel.ravel()[:-6]
palm_pos = self.data.site_xpos[self.S_grasp_sid].ravel()
sensordata = []
if use_tactile:
sensordata = self.data.sensordata.ravel().copy()[20:41]

res = np.concatenate([robot_jnt, robot_vel, palm_pos, sensordata])
return res

def reset_model(self):
qp = self.init_qpos.copy()
qv = self.init_qvel.copy()
Expand Down
12 changes: 11 additions & 1 deletion mj_envs/hand_manipulation_suite/relocate_v0.py
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,17 @@ def _get_obs(self):
palm_pos = self.data.site_xpos[self.S_grasp_sid].ravel()
target_pos = self.data.site_xpos[self.target_obj_sid].ravel()
return np.concatenate([qp[:-6], palm_pos-obj_pos, palm_pos-target_pos, obj_pos-target_pos])


def get_proprioception(self, use_tactile):
# return self._get_obs()
robot_pos = self.data.qpos.ravel()[:-6]
palm_pos = self.data.site_xpos[self.S_grasp_sid].ravel()
sensordata = []
if use_tactile:
sensordata = self.data.sensordata.ravel().copy()[20:41]
res = np.concatenate([robot_pos, palm_pos, sensordata])
return res

def reset_model(self):
qp = self.init_qpos.copy()
qv = self.init_qvel.copy()
Expand Down
14 changes: 14 additions & 0 deletions mj_envs/mujoco_env.py
Original file line number Diff line number Diff line change
Expand Up @@ -81,7 +81,15 @@ def viewer_setup(self):
"""
pass

def get_proprioception(self, use_tactile):
vikashplus marked this conversation as resolved.
Show resolved Hide resolved
"""
For the VIL paper
"""
pass

# -----------------------------
def get_body_com(self, body_name):
Copy link
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These are environment specific utilities that you can just add to the env definition. Im not sure if there is any benefit of adding this to the mujoco_env base class.

@aravindr93 - what do you think?

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am not sure if this function is correct in the first place. This is returning the xpos of the body, which is the root position of the body in global coordinates, but need not be center of mass. @vikashplus, is this correct? Ideally, you would have to look at all the members of this body, and do a weighted average of their xpos.

If this function is critical, and we are actually using it (meaning you need the COM and not just the root position), then it makes sense to have it in mujoco_env.py, since it might be useful for many envs and future applications.

Copy link
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Here is how you can get the COM
mjtNum* xipos; // Cartesian position of body com (nbody x 3)

Seems like get_body_xpos is giving you the location of the root of the body, not necessarily the body_com

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we can remove this function for now.

return self.sim.data.get_body_xpos(body_name)

def _reset(self):
self.sim.reset()
Expand Down Expand Up @@ -129,6 +137,12 @@ def state_vector(self):
state.qpos.flat, state.qvel.flat])

# -----------------------------
def get_pixels(self, frame_size=(128, 128), camera_name=None, device_id=0):
pixels = self.sim.render(width=frame_size[0], height=frame_size[1],
mode='offscreen', camera_name=camera_name, device_id=device_id)

pixels = pixels[::-1, :, :]
return pixels

def visualize_policy(self, policy, horizon=1000, num_episodes=1, mode='exploration'):
self.mujoco_render_frames = True
Expand Down