From 78d48c43003efe5cca30b159c231da68fc48cc8f Mon Sep 17 00:00:00 2001
From: Fabian <fabian.otto@de.bosch.com>
Date: Wed, 6 Jul 2022 09:19:21 +0200
Subject: [PATCH] mp wrappers added

---
 alr_envs/alr/mujoco/ant_jump/mp_wrapper.py    | 23 +++++++++++++++++++
 .../alr/mujoco/walker_2d_jump/mp_wrapper.py   | 23 +++++++++++++++++++
 2 files changed, 46 insertions(+)
 create mode 100644 alr_envs/alr/mujoco/ant_jump/mp_wrapper.py
 create mode 100644 alr_envs/alr/mujoco/walker_2d_jump/mp_wrapper.py

diff --git a/alr_envs/alr/mujoco/ant_jump/mp_wrapper.py b/alr_envs/alr/mujoco/ant_jump/mp_wrapper.py
new file mode 100644
index 0000000..a481d6f
--- /dev/null
+++ b/alr_envs/alr/mujoco/ant_jump/mp_wrapper.py
@@ -0,0 +1,23 @@
+from typing import Union, Tuple
+
+import numpy as np
+
+from alr_envs.black_box.raw_interface_wrapper import RawInterfaceWrapper
+
+
+class MPWrapper(RawInterfaceWrapper):
+
+    @property
+    def context_mask(self):
+        return np.hstack([
+            [False] * 111,  # ant has 111 dimensional observation space !!
+            [True]  # goal height
+        ])
+
+    @property
+    def current_pos(self) -> Union[float, int, np.ndarray]:
+        return self.env.sim.data.qpos[7:15].copy()
+
+    @property
+    def current_vel(self) -> Union[float, int, np.ndarray, Tuple]:
+        return self.env.sim.data.qvel[6:14].copy()
diff --git a/alr_envs/alr/mujoco/walker_2d_jump/mp_wrapper.py b/alr_envs/alr/mujoco/walker_2d_jump/mp_wrapper.py
new file mode 100644
index 0000000..63432a7
--- /dev/null
+++ b/alr_envs/alr/mujoco/walker_2d_jump/mp_wrapper.py
@@ -0,0 +1,23 @@
+from typing import Tuple, Union
+
+import numpy as np
+
+from alr_envs.black_box.raw_interface_wrapper import RawInterfaceWrapper
+
+
+class MPWrapper(RawInterfaceWrapper):
+
+    @property
+    def context_mask(self):
+        return np.hstack([
+            [False] * 17,
+            [True]  # goal pos
+        ])
+
+    @property
+    def current_pos(self) -> Union[float, int, np.ndarray]:
+        return self.env.data.qpos[3:9].copy()
+
+    @property
+    def current_vel(self) -> Union[float, int, np.ndarray, Tuple]:
+        return self.env.data.qvel[3:9].copy()