fancy_gym/alr_envs/alr/mujoco/beerpong/new_mp_wrapper.py

42 lines
1.3 KiB
Python
Raw Normal View History

2022-06-29 09:37:18 +02:00
from typing import Union, Tuple
2022-04-28 09:05:28 +02:00
import numpy as np
2022-06-30 17:33:05 +02:00
from alr_envs.black_box.raw_interface_wrapper import RawInterfaceWrapper
2022-04-28 09:05:28 +02:00
2022-06-30 14:55:34 +02:00
class MPWrapper(RawInterfaceWrapper):
2022-05-03 19:51:54 +02:00
@property
2022-04-28 09:05:28 +02:00
def current_pos(self) -> Union[float, int, np.ndarray, Tuple]:
return self.env.sim.data.qpos[0:7].copy()
2022-05-03 19:51:54 +02:00
@property
2022-04-28 09:05:28 +02:00
def current_vel(self) -> Union[float, int, np.ndarray, Tuple]:
return self.env.sim.data.qvel[0:7].copy()
2022-06-29 09:37:18 +02:00
def get_context_mask(self):
2022-04-28 09:05:28 +02:00
return np.hstack([
[False] * 7, # cos
[False] * 7, # sin
2022-06-02 09:05:38 +02:00
[False] * 7, # joint velocities
[False] * 3, # cup_goal_diff_final
[False] * 3, # cup_goal_diff_top
2022-04-28 09:05:28 +02:00
[True] * 2, # xy position of cup
[False] # env steps
2022-06-29 09:37:18 +02:00
])
2022-05-03 19:51:54 +02:00
# 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
2022-05-05 18:50:20 +02:00
return action, None
2022-05-03 19:51:54 +02:00
else:
return action, None
2022-05-29 11:58:01 +02:00
def set_context(self, context):
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())