v1.1.4; Add ZMP margin metric

robogauge/tasks/simulator/sim_data.py

@@ -5,7 +5,7 @@ from dataclasses import dataclass
 class JointState:
     pos: np.ndarray     # [rad] shape (n_dof,)
     vel: np.ndarray     # [rad/s] shape (n_dof,)
-    torque: np.ndarray   # [N*m] shape (n_dof,)
+    torque: np.ndarray  # [N*m] shape (n_dof,)
     limits: np.ndarray  # [rad] shape (n_dof, 2), lower and upper limits
     names: list         # list of joint names
 
@@ -30,9 +30,46 @@ class RobotProprioception:
     base: BaseState
     imu: IMUState
 
+@dataclass
+class RigidBodyDynamics:
+    names: list                   # list of rigid body names
+    mass: np.ndarray              # [kg] shape (n_body,)
+    com_pos: np.ndarray           # [m] world frame, shape (n_body, 3)
+    com_lin_vel: np.ndarray       # [m/s] world frame, shape (n_body, 3)
+    com_lin_acc: np.ndarray       # [m/s^2] world frame, shape (n_body, 3)
+    ang_vel: np.ndarray           # [rad/s] world frame, shape (n_body, 3)
+    ang_acc: np.ndarray           # [rad/s^2] world frame, shape (n_body, 3)
+    inertia_world: np.ndarray     # [kg*m^2] world frame, shape (n_body, 3, 3)
+
+@dataclass
+class GroundContactState:
+    positions: np.ndarray         # [m] world frame, shape (n_contact, 3)
+    distances: np.ndarray         # [m] shape (n_contact,)
+    robot_geom_names: list        # list of robot geom names
+    other_geom_names: list        # list of non-robot geom names
+    robot_body_names: list        # list of robot body names
+    other_body_names: list        # list of non-robot body names
+
+@dataclass
+class DynamicsState:
+    gravity: np.ndarray
+    rigid_bodies: RigidBodyDynamics
+    contacts: GroundContactState
+    default_diagonal_foot_distance: float
+
+@dataclass
+class VisualState:
+    zmp_world_pos: np.ndarray = None
+    zmp_draw_enabled: bool = False
+    zmp_draw_size: float = 0.03
+    zmp_draw_height_offset: float = 0.02
+    zmp_draw_rgba: np.ndarray = None
+
 @dataclass
 class SimData:
     n_step: int
     sim_dt: float
     sim_time: float
     proprio: RobotProprioception
+    dynamics: DynamicsState
+    visual: VisualState