walker.py

"""
MJX-compatible Rodent walker.
This one uses an xml that already includes a floor
"""

import os
import re

from dm_control import composer
from dm_control import mjcf
from dm_control.composer.observation import observable
from dm_control.locomotion.walkers import base
from dm_control.locomotion.walkers import legacy_base
from dm_control.mujoco import wrapper as mj_wrapper
import numpy as np

_XML_PATH = os.path.join(os.path.dirname(__file__), "assets/rodent.xml")

_RAT_MOCAP_JOINTS = [
    "vertebra_1_extend",
    "vertebra_2_bend",
    "vertebra_3_twist",
    "vertebra_4_extend",
    "vertebra_5_bend",
    "vertebra_6_twist",
    "hip_L_supinate",
    "hip_L_abduct",
    "hip_L_extend",
    "knee_L",
    "ankle_L",
    "toe_L",
    "hip_R_supinate",
    "hip_R_abduct",
    "hip_R_extend",
    "knee_R",
    "ankle_R",
    "toe_R",
    "vertebra_C1_extend",
    "vertebra_C1_bend",
    "vertebra_C2_extend",
    "vertebra_C2_bend",
    "vertebra_C3_extend",
    "vertebra_C3_bend",
    "vertebra_C4_extend",
    "vertebra_C4_bend",
    "vertebra_C5_extend",
    "vertebra_C5_bend",
    "vertebra_C6_extend",
    "vertebra_C6_bend",
    "vertebra_C7_extend",
    "vertebra_C9_bend",
    "vertebra_C11_extend",
    "vertebra_C13_bend",
    "vertebra_C15_extend",
    "vertebra_C17_bend",
    "vertebra_C19_extend",
    "vertebra_C21_bend",
    "vertebra_C23_extend",
    "vertebra_C25_bend",
    "vertebra_C27_extend",
    "vertebra_C29_bend",
    "vertebra_cervical_5_extend",
    "vertebra_cervical_4_bend",
    "vertebra_cervical_3_twist",
    "vertebra_cervical_2_extend",
    "vertebra_cervical_1_bend",
    "vertebra_axis_twist",
    "vertebra_atlant_extend",
    "atlas",
    "mandible",
    "scapula_L_supinate",
    "scapula_L_abduct",
    "scapula_L_extend",
    "shoulder_L",
    "shoulder_sup_L",
    "elbow_L",
    "wrist_L",
    "finger_L",
    "scapula_R_supinate",
    "scapula_R_abduct",
    "scapula_R_extend",
    "shoulder_R",
    "shoulder_sup_R",
    "elbow_R",
    "wrist_R",
    "finger_R",
]


_UPRIGHT_POS = (0.0, 0.0, 0.0)
_UPRIGHT_QUAT = (1.0, 0.0, 0.0, 0.0)
_TORQUE_THRESHOLD = 60


class Rat(legacy_base.Walker):
    """A position-controlled rat with control range scaled to [-1, 1]."""

    def _build(
        self,
        params=None,
        name="walker",
        torque_actuators=False,
        foot_mods=False,
        initializer=None,
    ):
        self.params = params
        self._mjcf_root = mjcf.from_path(_XML_PATH)
        if name:
            self._mjcf_root.model = name

        self.body_sites = []
        super()._build(initializer=initializer)

        # modify actuators
        if torque_actuators:
            for actuator in self._mjcf_root.find_all("actuator"):
                actuator.gainprm = [actuator.forcerange[1]]
                del actuator.biastype
                del actuator.biasprm

        # modify ankle and toe limits
        if foot_mods:
            self._mjcf_root.find("default", "ankle").joint.range = [-0.1, 2.0]
            self._mjcf_root.find("default", "toe").joint.range = [-0.7, 0.87]

    @property
    def upright_pose(self):
        """Reset pose to upright position."""
        return base.WalkerPose(xpos=_UPRIGHT_POS, xquat=_UPRIGHT_QUAT)

    @property
    def mjcf_model(self):
        """Return the model root."""
        return self._mjcf_root

    @composer.cached_property
    def actuators(self):
        """Return all actuators."""
        return tuple(self._mjcf_root.find_all("actuator"))

    @composer.cached_property
    def root_body(self):
        """Return the body."""
        return self._mjcf_root.find("body", "torso")

    @composer.cached_property
    def pelvis_body(self):
        """Return the body."""
        return self._mjcf_root.find("body", "pelvis")

    @composer.cached_property
    def head(self):
        """Return the head."""
        return self._mjcf_root.find("body", "skull")

    @composer.cached_property
    def left_arm_root(self):
        """Return the left arm."""
        return self._mjcf_root.find("body", "scapula_L")

    @composer.cached_property
    def right_arm_root(self):
        """Return the right arm."""
        return self._mjcf_root.find("body", "scapula_R")

    @composer.cached_property
    def ground_contact_geoms(self):
        """Return ground contact geoms."""
        return tuple(
            self._mjcf_root.find("body", "foot_L").find_all("geom")
            + self._mjcf_root.find("body", "foot_R").find_all("geom")
            + self._mjcf_root.find("body", "hand_L").find_all("geom")
            + self._mjcf_root.find("body", "hand_R").find_all("geom")
            + self._mjcf_root.find("body", "vertebra_C1").find_all("geom")
        )

    @composer.cached_property
    def standing_height(self):
        """Return standing height."""
        return self.params["_STAND_HEIGHT"]

    @composer.cached_property
    def end_effectors(self):
        """Return end effectors."""
        return (
            self._mjcf_root.find("body", "lower_arm_R"),
            self._mjcf_root.find("body", "lower_arm_L"),
            self._mjcf_root.find("body", "foot_R"),
            self._mjcf_root.find("body", "foot_L"),
        )

    @composer.cached_property
    def observable_joints(self):
        """Return observable joints."""
        return tuple(
            actuator.joint
            for actuator in self.actuators  #  This lint is mistaken; pylint: disable=not-an-iterable
            if actuator.joint is not None
        )

    @composer.cached_property
    def observable_tendons(self):
        return self._mjcf_root.find_all("tendon")

    @composer.cached_property
    def mocap_joints(self):
        return tuple(self._mjcf_root.find("joint", name) for name in _RAT_MOCAP_JOINTS)

    @composer.cached_property
    def mocap_joint_order(self):
        return tuple(
            [jnt.name for jnt in self.mocap_joints]
        )  #  This lint is mistaken; pylint: disable=not-an-iterable

    @composer.cached_property
    def bodies(self):
        """Return all bodies."""
        return tuple(self._mjcf_root.find_all("body"))

    @composer.cached_property
    def mocap_tracking_bodies(self):
        """Return bodies for mocap comparison."""
        return tuple(
            body
            for body in self._mjcf_root.find_all("body")
            if not re.match(r"(vertebra|hand|toe)", body.name)
        )

    @composer.cached_property
    def primary_joints(self):
        """Return primary (non-vertebra) joints."""
        return tuple(
            jnt
            for jnt in self._mjcf_root.find_all("joint")
            if "vertebra" not in jnt.name
        )

    @composer.cached_property
    def vertebra_joints(self):
        """Return vertebra joints."""
        return tuple(
            jnt for jnt in self._mjcf_root.find_all("joint") if "vertebra" in jnt.name
        )

    @composer.cached_property
    def primary_joint_order(self):
        joint_names = self.mocap_joint_order
        primary_names = tuple(
            [jnt.name for jnt in self.primary_joints]
        )  # pylint: disable=not-an-iterable
        primary_order = []
        for nm in primary_names:
            primary_order.append(joint_names.index(nm))
        return primary_order

    @composer.cached_property
    def vertebra_joint_order(self):
        joint_names = self.mocap_joint_order
        vertebra_names = tuple(
            [jnt.name for jnt in self.vertebra_joints]
        )  # pylint: disable=not-an-iterable
        vertebra_order = []
        for nm in vertebra_names:
            vertebra_order.append(joint_names.index(nm))
        return vertebra_order

    @composer.cached_property
    def egocentric_camera(self):
        """Return the egocentric camera."""
        return self._mjcf_root.find("camera", "egocentric")

    @property
    def _xml_path(self):
        """Return the path to th model .xml file."""
        return self.params["_XML_PATH"]

    @composer.cached_property
    def joint_actuators(self):
        """Return all joint actuators."""
        return tuple([act for act in self._mjcf_root.find_all("actuator") if act.joint])

    @composer.cached_property
    def joint_actuators_range(self):
        act_joint_range = []
        for (
            act
        ) in (
            self.joint_actuators
        ):  #  This lint is mistaken; pylint: disable=not-an-iterable
            associated_joint = self._mjcf_root.find("joint", act.name)
            act_range = associated_joint.dclass.joint.range
            act_joint_range.append(act_range)
        return act_joint_range

    def pose_to_actuation(self, pose):
        # holds for joint actuators, find desired torque = 0
        # u_ref = [2 q_ref - (r_low + r_up) ]/(r_up - r_low)
        r_lower = np.array(
            [ajr[0] for ajr in self.joint_actuators_range]
        )  #  This lint is mistaken; pylint: disable=not-an-iterable
        r_upper = np.array(
            [ajr[1] for ajr in self.joint_actuators_range]
        )  #  This lint is mistaken; pylint: disable=not-an-iterable
        num_tendon_actuators = len(self.actuators) - len(self.joint_actuators)
        tendon_actions = np.zeros(num_tendon_actuators)
        return np.hstack(
            [
                tendon_actions,
                (2 * pose[self.joint_actuator_order] - (r_lower + r_upper))
                / (r_upper - r_lower),
            ]
        )

    @composer.cached_property
    def joint_actuator_order(self):
        joint_names = self.mocap_joint_order
        joint_actuator_names = tuple(
            [act.name for act in self.joint_actuators]
        )  #  This lint is mistaken; pylint: disable=not-an-iterable
        actuator_order = []
        for nm in joint_actuator_names:
            actuator_order.append(joint_names.index(nm))
        return actuator_order

    def _build_observables(self):
        return RodentObservables(self)


class RodentObservables(legacy_base.WalkerObservables):
    """Observables for the Rat."""

    @composer.observable
    def head_height(self):
        """Observe the head height."""
        return observable.MJCFFeature("xpos", self._entity.head)[2]

    @composer.observable
    def sensors_torque(self):
        """Observe the torque sensors."""
        return observable.MJCFFeature(
            "sensordata",
            self._entity.mjcf_model.sensor.torque,
            corruptor=lambda v, random_state: np.tanh(2 * v / _TORQUE_THRESHOLD),
        )

    @composer.observable
    def tendons_pos(self):
        return observable.MJCFFeature("length", self._entity.observable_tendons)

    @composer.observable
    def tendons_vel(self):
        return observable.MJCFFeature("velocity", self._entity.observable_tendons)

    @composer.observable
    def actuator_activation(self):
        """Observe the actuator activation."""
        model = self._entity.mjcf_model
        return observable.MJCFFeature("act", model.find_all("actuator"))

    @composer.observable
    def appendages_pos(self):
        """Equivalent to `end_effectors_pos` with head's position appended."""

        def relative_pos_in_egocentric_frame(physics):
            end_effectors_with_head = self._entity.end_effectors + (self._entity.head,)
            end_effector = physics.bind(end_effectors_with_head).xpos
            torso = physics.bind(self._entity.root_body).xpos
            xmat = np.reshape(physics.bind(self._entity.root_body).xmat, (3, 3))
            return np.reshape(np.dot(end_effector - torso, xmat), -1)

        return observable.Generic(relative_pos_in_egocentric_frame)

    @property
    def proprioception(self):
        """Return proprioceptive information."""
        return [
            self.joints_pos,
            self.joints_vel,
            self.tendons_pos,
            self.tendons_vel,
            self.actuator_activation,
            self.body_height,
            self.end_effectors_pos,
            self.appendages_pos,
            self.world_zaxis,
        ] + self._collect_from_attachments("proprioception")

    @composer.observable
    def egocentric_camera(self):
        """Observable of the egocentric camera."""

        if not hasattr(self, "_scene_options"):
            # Don't render this walker's geoms.
            self._scene_options = mj_wrapper.MjvOption()
            collision_geom_group = 2
            self._scene_options.geomgroup[collision_geom_group] = 0
            cosmetic_geom_group = 1
            self._scene_options.geomgroup[cosmetic_geom_group] = 0

        return observable.MJCFCamera(
            self._entity.egocentric_camera,
            width=64,
            height=64,
            scene_option=self._scene_options,
        )

    @composer.cached_property
    def mocap_to_observable_joint_order(self):
        mocap_to_obs = [self.mocap_joints.index(j) for j in self.observable_joints]
        return mocap_to_obs