working bp version

alr_envs/mp/controllers/base_controller.py

@@ -0,0 +1,4 @@
+class BaseController:
+
+    def get_action(self, des_pos, des_vel, c_pos, c_vel):
+        raise NotImplementedError

alr_envs/mp/controllers/controller_factory.py

@@ -0,0 +1,21 @@
+from alr_envs.mp.controllers.meta_world_controller import MetaWorldController
+from alr_envs.mp.controllers.pd_controller import PDController
+from alr_envs.mp.controllers.vel_controller import VelController
+from alr_envs.mp.controllers.pos_controller import PosController
+
+ALL_TYPES = ["motor", "velocity", "position", "metaworld"]
+
+
+def get_controller(controller_type: str, **kwargs):
+    controller_type = controller_type.lower()
+    if controller_type == "motor":
+        return PDController(**kwargs)
+    elif controller_type == "velocity":
+        return VelController()
+    elif controller_type == "position":
+        return PosController()
+    elif controller_type == "metaworld":
+        return MetaWorldController()
+    else:
+        raise ValueError(f"Specified controller type {controller_type} not supported, "
+                         f"please choose one of {ALL_TYPES}.")

alr_envs/mp/controllers/meta_world_controller.py

@@ -0,0 +1,25 @@
+import numpy as np
+
+from alr_envs.mp.controllers.base_controller import BaseController
+
+
+class MetaWorldController(BaseController):
+    """
+    A Metaworld Controller. Using position and velocity information from a provided environment,
+    the controller calculates a response based on the desired position and velocity.
+    Unlike the other Controllers, this is a special controller for MetaWorld environments.
+    They use a position delta for the xyz coordinates and a raw position for the gripper opening.
+
+    :param env: A position environment
+    """
+
+    def get_action(self, des_pos, des_vel, c_pos, c_vel):
+        gripper_pos = des_pos[-1]
+
+        cur_pos = env.current_pos[:-1]
+        xyz_pos = des_pos[:-1]
+
+        assert xyz_pos.shape == cur_pos.shape, \
+            f"Mismatch in dimension between desired position {xyz_pos.shape} and current position {cur_pos.shape}"
+        trq = np.hstack([(xyz_pos - cur_pos), gripper_pos])
+        return trq

alr_envs/mp/controllers/pd_controller.py

@@ -0,0 +1,28 @@
+from typing import Union, Tuple
+
+from alr_envs.mp.controllers.base_controller import BaseController
+
+
+class PDController(BaseController):
+    """
+    A PD-Controller. Using position and velocity information from a provided environment,
+    the controller calculates a response based on the desired position and velocity
+
+    :param env: A position environment
+    :param p_gains: Factors for the proportional gains
+    :param d_gains: Factors for the differential gains
+    """
+
+    def __init__(self,
+                 p_gains: Union[float, Tuple] = 1,
+                 d_gains: Union[float, Tuple] = 0.5):
+        self.p_gains = p_gains
+        self.d_gains = d_gains
+
+    def get_action(self, des_pos, des_vel, c_pos, c_vel):
+        assert des_pos.shape == c_pos.shape, \
+            f"Mismatch in dimension between desired position {des_pos.shape} and current position {c_pos.shape}"
+        assert des_vel.shape == c_vel.shape, \
+            f"Mismatch in dimension between desired velocity {des_vel.shape} and current velocity {c_vel.shape}"
+        trq = self.p_gains * (des_pos - c_pos) + self.d_gains * (des_vel - c_vel)
+        return trq

alr_envs/mp/controllers/pos_controller.py

@@ -0,0 +1,9 @@
+from alr_envs.mp.controllers.base_controller import BaseController
+
+
+class PosController(BaseController):
+    """
+    A Position Controller. The controller calculates a response only based on the desired position.
+    """
+    def get_action(self, des_pos, des_vel, c_pos, c_vel):
+        return des_pos

alr_envs/mp/controllers/vel_controller.py

@@ -0,0 +1,9 @@
+from alr_envs.mp.controllers.base_controller import BaseController
+
+
+class VelController(BaseController):
+    """
+    A Velocity Controller. The controller calculates a response only based on the desired velocity.
+    """
+    def get_action(self, des_pos, des_vel, c_pos, c_vel):
+        return des_vel