v0.1.3

robogauge/tasks/robots/base_robot.py

@@ -9,21 +9,55 @@
 '''
 import torch
 import numpy as np
+
+from robogauge.utils.helpers import parse_path
+from robogauge.utils.logger import logger
 from robogauge.tasks.robots.base_robot_config import RobotConfig
+from robogauge.tasks.simulator.sim_data import SimData
+from robogauge.tasks.gauge.goal_data import GoalData
 
 class BaseRobot:
     def __init__(self, cfg: RobotConfig):
-        self.num_act = cfg.mdp.num_actions
-        self.num_obs = cfg.mdp.num_observations
-        self.model = None
+        self.cfg = cfg
+        self.device = self.cfg.control.device
+        self.num_obs = cfg.control.num_observations
+        self.num_action = cfg.control.num_actions
+        self.control_type = cfg.control.control_type
+        self.p_gains = np.array(cfg.control.p_gains)
+        self.d_gains = np.array(cfg.control.d_gains)
+        script_model_path = parse_path(cfg.control.torch_script_model_path)
+        logger.info(f"Loading robot model from '{script_model_path}'")
+        self.model = torch.jit.load(script_model_path).to(self.device)
+        self.model.eval()
     
-    def load_model(self):
-        ...
-    
-    def build_observation(self, sim_info: dict, goal_info: dict) -> np.ndarray:
-        obs = np.zeros(self.num_obs)
+    def build_observation(self, sim_data: SimData, goal_data: GoalData) -> np.ndarray:
+        obs = np.zeros(self.num_obs, dtype=np.float32)
         return obs
     
-    def get_action(self, obs) -> np.ndarray:
-        action = np.zeros_like(self.num_act)
-        return action
+    def get_action(self, obs: np.ndarray):
+        """
+        Returns:
+            action: (num_action,) target joint positions/velocities/torques
+            p_gains: (num_action,) proportional gains for Mujoco PD controller
+            d_gains: (num_action,) derivative gains for Mujoco PD controller
+            control_type: 'P', 'V', or 'T' for position/velocity/torque control
+        """
+        action = np.zeros(self.num_action, dtype=np.float32)
+        return action, self.p_gains, self.d_gains, self.control_type
+
+def get_projected_gravity(quat):
+    """ Compute world frame gravity (0, 0, -1) projected into robot base frame.
+    Args:
+        quat: (4,) quaternion (w, x, y, z) from robot base to world frame
+    Returns:
+        projected_gravity: (3,) projected gravity vector in robot base frame
+    """
+    qw, qx, qy, qz = quat
+
+    gravity_orientation = np.zeros(3)
+
+    gravity_orientation[0] = 2 * (-qz * qx + qw * qy)
+    gravity_orientation[1] = -2 * (qz * qy + qw * qx)
+    gravity_orientation[2] = 1 - 2 * (qw * qw + qz * qz)
+
+    return gravity_orientation