Source code for wbc_mjlab.presets.se_actor

"""State-estimation actor observation preset.

SE tasks drop height/gravity reference proxies and add anchor pose tracking error
(world position + axis-angle orientation error) plus ``base_lin_vel``.
"""

from __future__ import annotations

from mjlab.envs import ManagerBasedRlEnvCfg
from mjlab.envs.mdp import builtin_sensor
from mjlab.managers.observation_manager import ObservationTermCfg
from mjlab.utils.noise import UniformNoiseCfg as Unoise

from wbc_mjlab.env.mdp.observations import (
  motion_anchor_ori_error,
  motion_anchor_pos_error_w,
)

_MOTION = "motion"
_MOTION_PARAMS = {"command_name": _MOTION}

_SE_REMOVED_TERMS = (
  "ref_base_height",
  "ref_gravity_b",
  "projected_gravity",
)

# Anchor (``MotionCommand.anchor_*`` vs ``robot_anchor_*``), not pelvis root — matches
# tracking rewards/terminations; ``base_lin_vel`` remains root-frame IMU velocity.
_SE_ADDED_TERMS: dict[str, ObservationTermCfg] = {
  "motion_anchor_pos_error_w": ObservationTermCfg(
    func=motion_anchor_pos_error_w,
    params=_MOTION_PARAMS,
    noise=Unoise(n_min=-0.1, n_max=0.1),
  ),
  "motion_anchor_ori_error": ObservationTermCfg(
    func=motion_anchor_ori_error,
    params=_MOTION_PARAMS,
    noise=Unoise(n_min=-0.02, n_max=0.02),
  ),
  "base_lin_vel": ObservationTermCfg(
    func=builtin_sensor,
    params={"sensor_name": ""},
    noise=Unoise(n_min=-0.2, n_max=0.2),
  ),
}




def apply_se_actor(cfg: ManagerBasedRlEnvCfg) -> None:
  """State-estimation actor layout: drop height/gravity proxies; add anchor pose error + base lin vel."""
  actor = cfg.observations["actor"]
  for key in _SE_REMOVED_TERMS:
    actor.terms.pop(key, None)
  actor.terms.update(_SE_ADDED_TERMS)