dodox/fancy_gym
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

removed dependency on mujoco-py

Browse Source
This commit is contained in:
philipp
2022-07-13 13:28:39 +02:00
parent 2fbf42167f
commit b04e7c1dcd
17 changed files with 134 additions and 80 deletions
Download Patch File Download Diff File Expand all files Collapse all files
alr_envs/envs/mujoco/beerpong/assets/beerpong_wo_cup_big_table.xml
+5
View File
@@ -185,4 +185,9 @@
<motor ctrllimited="true" ctrlrange="-1.0 1.0" gear="2.0" joint="wam/palm_yaw_joint"/>
</actuator>
<sensor>
<framelinvel name="init_ball_vel" objtype="site" objname="init_ball_pos" />
</sensor>
</mujoco>
alr_envs/envs/mujoco/beerpong/beerpong.py
+19 -20
View File
@@ -1,7 +1,6 @@
import os
from typing import Optional
import mujoco_py.builder
import numpy as np
from gym import utils
from gym.envs.mujoco import MujocoEnv
@@ -50,9 +49,9 @@ class BeerPongEnv(MujocoEnv, utils.EzPickle):
self.repeat_action = frame_skip
# TODO: If accessing IDs is easier in the (new) official mujoco bindings, remove this
self.model = None
self.site_id = lambda x: self.model.site_name2id(x)
self.body_id = lambda x: self.model.body_name2id(x)
self.geom_id = lambda x: self.model.geom_name2id(x)
self.geom_id = lambda x: self._mujoco_bindings.mj_name2id(self.model,
self._mujoco_bindings.mjtObj.mjOBJ_GEOM,
x)
# for reward calculation
self.dists = []
@@ -65,7 +64,7 @@ class BeerPongEnv(MujocoEnv, utils.EzPickle):
self.ball_in_cup = False
self.dist_ground_cup = -1 # distance floor to cup if first floor contact
MujocoEnv.__init__(self, self.xml_path, frame_skip=1, mujoco_bindings="mujoco_py")
MujocoEnv.__init__(self, model_path=self.xml_path, frame_skip=1, mujoco_bindings="mujoco")
utils.EzPickle.__init__(self)
@property
@@ -100,13 +99,13 @@ class BeerPongEnv(MujocoEnv, utils.EzPickle):
# TODO: Ask Max why we need to set the state twice.
self.set_state(start_pos, init_vel)
start_pos[7::] = self.data.site_xpos[self.site_id("init_ball_pos"), :].copy()
start_pos[7::] = self.data.site("init_ball_pos").xpos.copy()
self.set_state(start_pos, init_vel)
xy = self.np_random.uniform(self._cup_pos_min, self._cup_pos_max)
xyz = np.zeros(3)
xyz[:2] = xy
xyz[-1] = 0.840
self.model.body_pos[self.body_id("cup_table")] = xyz
self.model.body("cup_table").pos[:] = xyz
return self._get_obs()
def step(self, a):
@@ -117,10 +116,10 @@ class BeerPongEnv(MujocoEnv, utils.EzPickle):
self.do_simulation(applied_action, self.frame_skip)
# self.reward_function.check_contacts(self.sim) # I assume this is not important?
if self._steps < self.release_step:
self.data.qpos[7::] = self.data.site_xpos[self.site_id("init_ball_pos"), :].copy()
self.data.qvel[7::] = self.data.site_xvelp[self.site_id("init_ball_pos"), :].copy()
self.data.qpos[7::] = self.data.site('init_ball_pos').xpos.copy()
self.data.qvel[7::] = self.data.sensor('init_ball_vel').data.copy()
crash = False
except mujoco_py.builder.MujocoException:
except Exception as e:
crash = True
ob = self._get_obs()
@@ -145,18 +144,18 @@ class BeerPongEnv(MujocoEnv, utils.EzPickle):
return ob, reward, done, infos
def _get_obs(self):
theta = self.data.qpos.flat[:7]
theta_dot = self.data.qvel.flat[:7]
ball_pos = self.data.get_body_xpos("ball").copy()
cup_goal_diff_final = ball_pos - self.data.get_site_xpos("cup_goal_final_table").copy()
cup_goal_diff_top = ball_pos - self.data.get_site_xpos("cup_goal_table").copy()
theta = self.data.qpos.flat[:7].copy()
theta_dot = self.data.qvel.flat[:7].copy()
ball_pos = self.data.qpos.flat[7:].copy()
cup_goal_diff_final = ball_pos - self.data.site("cup_goal_final_table").xpos.copy()
cup_goal_diff_top = ball_pos - self.data.site("cup_goal_table").xpos.copy()
return np.concatenate([
np.cos(theta),
np.sin(theta),
theta_dot,
cup_goal_diff_final,
cup_goal_diff_top,
self.model.body_pos[self.body_id("cup_table")][:2].copy(),
self.model.body("cup_table").pos[:2].copy(),
# [self._steps], # Use TimeAwareObservation Wrapper instead ....
])
@@ -165,10 +164,10 @@ class BeerPongEnv(MujocoEnv, utils.EzPickle):
return super(BeerPongEnv, self).dt * self.repeat_action
def _get_reward(self, action):
goal_pos = self.data.get_site_xpos("cup_goal_table")
ball_pos = self.data.get_body_xpos("ball")
ball_vel = self.data.get_body_xvelp("ball")
goal_final_pos = self.data.get_site_xpos("cup_goal_final_table")
goal_pos = self.data.site("cup_goal_table").xpos
goal_final_pos = self.data.site("cup_goal_final_table").xpos
ball_pos = self.data.qpos.flat[7:].copy()
ball_vel = self.data.qvel.flat[7:].copy()
self._check_contacts()
self.dists.append(np.linalg.norm(goal_pos - ball_pos))
alr_envs/envs/mujoco/beerpong/mp_wrapper.py
+5 -5
View File
@@ -21,16 +21,16 @@ class MPWrapper(RawInterfaceWrapper):
@property
def current_pos(self) -> Union[float, int, np.ndarray, Tuple]:
return self.env.data.qpos[0:7].copy()
return self.data.qpos[0:7].copy()
@property
def current_vel(self) -> Union[float, int, np.ndarray, Tuple]:
return self.env.data.qvel[0:7].copy()
return self.data.qvel[0:7].copy()
# TODO: Fix this
def _episode_callback(self, action: np.ndarray, mp) -> Tuple[np.ndarray, Union[np.ndarray, None]]:
if mp.learn_tau:
self.env.env.release_step = action[0] / self.env.dt # Tau value
self.release_step = action[0] / self.dt # Tau value
return action, None
else:
return action, None
@@ -39,5 +39,5 @@ class MPWrapper(RawInterfaceWrapper):
xyz = np.zeros(3)
xyz[:2] = context
xyz[-1] = 0.840
self.env.env.model.body_pos[self.env.env.cup_table_id] = xyz
return self.get_observation_from_step(self.env.env._get_obs())
self.model.body_pos[self.cup_table_id] = xyz
return self.get_observation_from_step(self.get_obs())