Skip to content

Motion Optimization

Offline optimization of motion clips before training. The one optimizer that ships is TCRS, which turns flat-ground biped clips into terrain-adapted ones: gait-phase and contact detection, MPPI foot-swing planning against a terrain height field, support-aware pelvis-height filtering, and whole-body Jacobian IK on a MuJoCo model of the robot. It needs no Eden env (numpy, scipy and mujoco only). Datasets reach it through the TCRSTerrainAdapt dataset modifier. The TCRS terrain adaptation guide explains the method and its settings.

from eden.extensions.motion_optimization.tcrs import HeightFieldTerrain, TCRSPlannerOptions, TCRSRobotOptions

Options

Configuration options for the offline motion-optimization extensions.

Holds the config surface of the TCRS terrain-adaptive motion optimizer (:mod:eden.extensions.motion_optimization.tcrs).

Upstream TCRS hardcodes Unitree G1 geometry (body/joint names, foot contact offsets, leg reach) across several modules. The port lifts all of it into :class:TCRSRobotOptions so another biped can be targeted without touching the algorithms; the no-argument default reproduces the upstream G1 values with the name lists taken from Eden's canonical G1 definitions.

:class:TCRSPlannerOptions collects the pipeline knobs that upstream exposed as CLI flags (defaults preserved). Both are :class:ConfigurableOptions (extra="forbid"), matching Eden's library-wide config pattern; the TCRS engine consumes them as read-only attribute bags and stays env-free.

Attributes

DEFAULT_IK_CONFIG module-attribute

DEFAULT_IK_CONFIG: dict = {'link_weights': {'pelvis': {'pos_weight': 1, 'rot_weight': 1.0}, 'left_hip_pitch_link': {'pos_weight': 0.0, 'rot_weight': 0.0}, 'left_hip_roll_link': {'pos_weight': 0.0, 'rot_weight': 0.0}, 'left_hip_yaw_link': {'pos_weight': 0.0, 'rot_weight': 0.0}, 'left_knee_link': {'pos_weight': 0.0, 'rot_weight': 0.0}, 'left_ankle_pitch_link': {'pos_weight': 0.0, 'rot_weight': 0.0}, 'right_hip_pitch_link': {'pos_weight': 0.0, 'rot_weight': 0.0}, 'right_hip_roll_link': {'pos_weight': 0.0, 'rot_weight': 0.0}, 'right_hip_yaw_link': {'pos_weight': 0.0, 'rot_weight': 0.0}, 'right_knee_link': {'pos_weight': 0.0, 'rot_weight': 0.0}, 'right_ankle_pitch_link': {'pos_weight': 0.0, 'rot_weight': 0.0}, 'left_ankle_roll_link': {'pos_weight': 10.0, 'rot_weight': 6.0}, 'right_ankle_roll_link': {'pos_weight': 10.0, 'rot_weight': 6.0}, 'left_toe': {'pos_weight': 6.0, 'rot_weight': 0.0, 'offset': [0.12, 0.0, -0.03]}, 'left_heel': {'pos_weight': 6.0, 'rot_weight': 0.0, 'offset': [-0.05, 0.0, -0.03]}, 'right_toe': {'pos_weight': 6.0, 'rot_weight': 0.0, 'offset': [0.12, 0.0, -0.03]}, 'right_heel': {'pos_weight': 6.0, 'rot_weight': 0.0, 'offset': [-0.05, 0.0, -0.03]}}, 'solver': {'max_iters': 60, 'tol': 0.001, 'damping': 0.05, 'step_size': 0.5, 'orientation_weight': 0.2, 'posture_weight': 0.01, 'max_joint_step': 0.2, 'penetration_weight': 20.0, 'root_xy_weight': 3.0, 'root_z_weight': 20, 'knee_forward_weight': 15.0, 'multiseed_enabled': True, 'multiseed_num_seeds': 3, 'multiseed_trigger_pos_err': 0.2, 'multiseed_trigger_joint_step': 0.8, 'multiseed_seed_span': 0.35, 'multiseed_continuity_weight': 0.25, 'multiseed_posture_weight': 0.02}}

Classes

TCRSRobotOptions

TCRSRobotOptions(**data)

Bases: ConfigurableOptions

Robot geometry and naming the TCRS pipeline needs.

The no-argument default is the upstream Unitree G1 spec with Eden's canonical name lists as clip orderings. There is no separate override API: construction is always overrides-on-the-G1-defaults — any unspecified field keeps its G1 value, and targeting another biped just means overriding more of them. Eden clips store joints in G1_DOFS29_NAME (chain) order; upstream PMT clips use an interleaved "policy" order. Both are supported — pass the clip's actual ordering; every consumer resolves columns by name.

Source code in eden/options/options.py
def __init__(self, **data):
    # Seed type metadata as construction data so BaseModel validates it into
    # __pydantic_extra__ (no-op for extra="forbid" classes, which use
    # the serialization envelope instead).
    data.update(option_metadata(self))
    # NOTE: all of the option modifications/validations should be done in the __init__ method
    try:
        BaseModel.__init__(self, **data)
    except ValidationError as exc:
        # A misspelled field on an extra="forbid" class should teach the field set, not just
        # name the typo. Only rewritten when every error is extra_forbidden.
        enriched = _enrich_extra_forbidden(type(self), exc)
        if enriched is None:
            raise
        raise enriched from None

Attributes

dofs_name class-attribute instance-attribute
dofs_name: list[str] = Field(default_factory=lambda: list(G1_DOFS29_NAME))

DOF order of the motion clip's dofs_pos columns. All MuJoCo qpos addressing is resolved by these names — never by column index.

links_name: list[str] = Field(default_factory=lambda: list(G1_DOFS29_TRACKED_LINKS_NAME))

Link order of the motion clip's links_pos_w/links_quat_w rows (pelvis first).

pelvis_link: str = 'pelvis'
left_foot_link: str = 'left_ankle_roll_link'
right_foot_link: str = 'right_ankle_roll_link'
left_hip_link: str = 'left_hip_pitch_link'
right_hip_link: str = 'right_hip_pitch_link'
left_knee_link: str = 'left_knee_link'
right_knee_link: str = 'right_knee_link'
left_leg_dofs_name class-attribute instance-attribute
left_leg_dofs_name: tuple[str, ...] = ('left_hip_pitch_joint', 'left_hip_roll_joint', 'left_hip_yaw_joint', 'left_knee_joint', 'left_ankle_pitch_joint', 'left_ankle_roll_joint')
right_leg_dofs_name class-attribute instance-attribute
right_leg_dofs_name: tuple[str, ...] = ('right_hip_pitch_joint', 'right_hip_roll_joint', 'right_hip_yaw_joint', 'right_knee_joint', 'right_ankle_pitch_joint', 'right_ankle_roll_joint')
waist_dofs_name class-attribute instance-attribute
waist_dofs_name: tuple[str, ...] = ('waist_yaw_joint', 'waist_roll_joint', 'waist_pitch_joint')
toe_offset class-attribute instance-attribute
toe_offset: tuple[float, float, float] = (0.12, 0.0, -0.03)

Representative toe contact point in the ankle-roll-link frame (from the G1 MJCF contact spheres).

heel_offset class-attribute instance-attribute
heel_offset: tuple[float, float, float] = (-0.05, 0.0, -0.03)

Representative heel contact point in the ankle-roll-link frame.

toe_probe_offset_x class-attribute instance-attribute
toe_probe_offset_x: float = 0.14

Forward probe distance used by footstep resolution (upstream TOE_OFFSET_X).

foot_nominal_z class-attribute instance-attribute
foot_nominal_z: float = 0.06

Ankle height above the support surface when the foot is flat on the ground.

hip_offset_z class-attribute instance-attribute
hip_offset_z: float = -0.251

Pelvis origin to hip joint, z component (m).

leg_max_length class-attribute instance-attribute
leg_max_length: float = 0.494

Hip-to-ankle chain length at full extension (m).

pelvis_ankle_reach_slack class-attribute instance-attribute
pelvis_ankle_reach_slack: float = 0.06

Extra reach beyond |hip_offset_z| + leg_max_length from the hip chain's lateral/forward offsets; the pipeline works in pelvis-link ↔ ankle-link coordinates.

ik_config class-attribute instance-attribute
ik_config: dict = Field(default_factory=dict)

Per-link IK weights and solver settings (schema of ik_config.json). Partial dicts — including this empty default — deep-merge onto :data:DEFAULT_IK_CONFIG at construction (see :func:merged_ik_config), so the stored value is always the full config.

pelvis_index property
pelvis_index: int
left_foot_index property
left_foot_index: int
right_foot_index property
right_foot_index: int
left_hip_index property
left_hip_index: int
right_hip_index property
right_hip_index: int
left_knee_index property
left_knee_index: int
right_knee_index property
right_knee_index: int
mid_foot_offset property
mid_foot_offset: tuple[float, float, float]

Midpoint of toe and heel contact points (upstream MID_FOOT_OFFSET).

pelvis_ankle_reach property
pelvis_ankle_reach: float
max_reach property
max_reach: float

Reach cap used when shaping pelvis height (upstream _MAX_REACH).

TCRSPlannerOptions

TCRSPlannerOptions(**data)

Bases: ConfigurableOptions

TCRS pipeline knobs, mirroring the upstream CLI defaults.

Not ported (dropped features): planner="plum" (imports a module that does not exist upstream), ik_backend curobo/drake, ik_mode="window" trajectory IK, and rolling_interval viewer-time replanning.

Range constraints live on the fields (gt/ge/le/allow_inf_nan), mirroring the upstream valid ranges.

Source code in eden/options/options.py
def __init__(self, **data):
    # Seed type metadata as construction data so BaseModel validates it into
    # __pydantic_extra__ (no-op for extra="forbid" classes, which use
    # the serialization envelope instead).
    data.update(option_metadata(self))
    # NOTE: all of the option modifications/validations should be done in the __init__ method
    try:
        BaseModel.__init__(self, **data)
    except ValidationError as exc:
        # A misspelled field on an extra="forbid" class should teach the field set, not just
        # name the typo. Only rewritten when every error is extra_forbidden.
        enriched = _enrich_extra_forbidden(type(self), exc)
        if enriched is None:
            raise
        raise enriched from None

Attributes

contact_z_threshold class-attribute instance-attribute
contact_z_threshold: float = Field(0.06, ge=0)
contact_speed_threshold class-attribute instance-attribute
contact_speed_threshold: float = Field(0.5, gt=0)
lookahead class-attribute instance-attribute
lookahead: float = Field(0.0, ge=0)
smoothing_alpha class-attribute instance-attribute
smoothing_alpha: float = Field(0.08, gt=0, le=1)
heel_offset_x class-attribute instance-attribute
heel_offset_x: float = -0.06
mid_offset_x class-attribute instance-attribute
mid_offset_x: float = 0.15
swing_floor_margin class-attribute instance-attribute
swing_floor_margin: float = 0.003
w_track class-attribute instance-attribute
w_track: float = Field(10.0, allow_inf_nan=False)
w_terrain class-attribute instance-attribute
w_terrain: float = Field(1000.0, allow_inf_nan=False)
h_clearance class-attribute instance-attribute
h_clearance: float = Field(0.03, ge=0)
mppi_n_knots class-attribute instance-attribute
mppi_n_knots: int = Field(5, gt=0)
mppi_n_samples class-attribute instance-attribute
mppi_n_samples: int = Field(128, gt=0)
mppi_n_iterations class-attribute instance-attribute
mppi_n_iterations: int = Field(10, gt=0)
mppi_temperature class-attribute instance-attribute
mppi_temperature: float = Field(0.1, gt=0)
leg_collision_iters class-attribute instance-attribute
leg_collision_iters: int = Field(4, ge=0)
leg_collision_margin class-attribute instance-attribute
leg_collision_margin: float = 0.025
shin_collision_radius class-attribute instance-attribute
shin_collision_radius: float = 0.06
foot_collision_radius class-attribute instance-attribute
foot_collision_radius: float = 0.04
collision_push_gain class-attribute instance-attribute
collision_push_gain: float = 0.9
collision_max_push class-attribute instance-attribute
collision_max_push: float = 0.06
pelvis_side_margin class-attribute instance-attribute
pelvis_side_margin: float = 0.04
pelvis_side_gain class-attribute instance-attribute
pelvis_side_gain: float = 0.5
stance_latch class-attribute instance-attribute
stance_latch: bool = True
stance_latch_ground_tol class-attribute instance-attribute
stance_latch_ground_tol: float = 0.05
ik_lower_body_stabilizer class-attribute instance-attribute
ik_lower_body_stabilizer: bool = True
ik_stabilizer_stance_scale class-attribute instance-attribute
ik_stabilizer_stance_scale: float = 2.2
ik_stabilizer_swing_scale class-attribute instance-attribute
ik_stabilizer_swing_scale: float = 0.6
ik_stabilizer_root_xy_scale class-attribute instance-attribute
ik_stabilizer_root_xy_scale: float = 0.6
ik_stabilizer_root_z_scale class-attribute instance-attribute
ik_stabilizer_root_z_scale: float = 0.75
ik_stabilizer_root_continuity_weight class-attribute instance-attribute
ik_stabilizer_root_continuity_weight: float = 80.0
ik_stabilizer_trigger_mm class-attribute instance-attribute
ik_stabilizer_trigger_mm: float = 60.0
ik_stabilizer_cooldown_frames class-attribute instance-attribute
ik_stabilizer_cooldown_frames: int = 6
lowpass_cutoff_hz class-attribute instance-attribute
lowpass_cutoff_hz: float = Field(5.0, ge=0)
lowpass_order class-attribute instance-attribute
lowpass_order: int = Field(4, gt=0)
reject_ik_mm class-attribute instance-attribute
reject_ik_mm: float = 200.0
seed class-attribute instance-attribute
seed: int = 42

Functions:

merged_ik_solver_config

merged_ik_solver_config(solver_overrides: dict) -> dict

Merge a (possibly partial) ik_config["solver"] section over the defaults.

DEFAULT_IK_CONFIG["solver"] is the complete set of solver knobs; every consumer (:meth:~eden.extensions.motion_optimization.tcrs.ik.JacobianIK.from_config, the pipeline's root-weight lookup) reads the merged dict, so a partial override changes exactly the named keys and never falls back to a second, disagreeing set of hardcoded defaults. Unknown keys raise.

Source code in eden/options/extensions/motion_optimization.py
def merged_ik_solver_config(solver_overrides: dict) -> dict:
    """Merge a (possibly partial) ``ik_config["solver"]`` section over the defaults.

    ``DEFAULT_IK_CONFIG["solver"]`` is the complete set of solver knobs; every consumer
    (:meth:`~eden.extensions.motion_optimization.tcrs.ik.JacobianIK.from_config`, the pipeline's
    root-weight lookup) reads the merged dict, so a partial
    override changes exactly the named keys and never falls back to a second,
    disagreeing set of hardcoded defaults. Unknown keys raise.
    """
    valid = set(DEFAULT_IK_CONFIG["solver"])
    unknown = set(solver_overrides) - valid
    if unknown:
        raise ValueError(f"Unknown ik_config solver key(s) {sorted(unknown)}. Valid solver keys: {sorted(valid)}.")
    return {**copy.deepcopy(DEFAULT_IK_CONFIG["solver"]), **solver_overrides}

merged_ik_config

merged_ik_config(overrides: dict) -> dict

Deep-merge a partial ik_config dict onto :data:DEFAULT_IK_CONFIG.

Merge semantics:

  • solver: key-wise merge over the default solver section; unknown solver keys raise (see :func:merged_ik_solver_config).
  • link_weights: per-link-name merge — each override entry is merged over the default entry of the same name (so {"pelvis": {"pos_weight": 2}} keeps the default rot_weight); names absent from the defaults are added as new constraints.

Sections other than solver/link_weights raise.

Source code in eden/options/extensions/motion_optimization.py
def merged_ik_config(overrides: dict) -> dict:
    """Deep-merge a partial ``ik_config`` dict onto :data:`DEFAULT_IK_CONFIG`.

    Merge semantics:

    - ``solver``: key-wise merge over the default solver section; unknown
      solver keys raise (see :func:`merged_ik_solver_config`).
    - ``link_weights``: per-link-name merge — each override entry is merged
      over the default entry of the same name (so ``{"pelvis": {"pos_weight": 2}}``
      keeps the default ``rot_weight``); names absent from the defaults are
      added as new constraints.

    Sections other than ``solver``/``link_weights`` raise.
    """
    if not isinstance(overrides, dict):
        raise ValueError(f"ik_config overrides must be a dict; got {type(overrides).__name__}.")
    if "body_weights" in overrides:
        raise ValueError(
            "ik_config has no 'body_weights' section: Eden names it 'link_weights' (its dofs_*/links_* naming "
            "convention); rename the section to 'link_weights'."
        )
    unknown_sections = set(overrides) - {"link_weights", "solver"}
    if unknown_sections:
        raise ValueError(
            f"Unknown ik_config section(s) {sorted(unknown_sections)}. Valid sections: ['link_weights', 'solver']."
        )
    merged = copy.deepcopy(DEFAULT_IK_CONFIG)
    merged["solver"] = merged_ik_solver_config(overrides.get("solver", {}))
    for link_name, weights in overrides.get("link_weights", {}).items():
        merged["link_weights"][link_name] = {**merged["link_weights"].get(link_name, {}), **weights}
    return merged

Pipeline

Offline TCRS pipeline: terrain-adaptive whole-body motion optimization.

Pipeline stages (per clip): 1. Build a :class:~eden.extensions.motion_optimization.tcrs.motion.MotionClip from raw arrays (contact masks + gait-phase clock are derived on construction). 2. Apply an SE(2) placement transform (yaw about the frame-0 pelvis, then xy translation). 3. Build :class:~eden.extensions.motion_optimization.tcrs.motion.TerrainReference + FootstepResolver (snap footholds onto safe terrain support surfaces). 4. MPPI-optimize every swing-foot trajectory on the mid-foot virtual point (:func:precompute_mppi_foot), then post-process: lateral separation clamps, side-wall pushes, capsule-based leg/foot collision resolution, stance latching. 5. Filter the pelvis z with :class:~eden.extensions.motion_optimization.tcrs.terrain_warp.SupportAwareRootZFilter against a reach-bounded terrain-driven target. 6. Solve per-frame whole-body Jacobian IK (with the lower-body stance stabilizer post-pass) and record raw / ghost-FK / optimized motion sets. 7. Zero-phase lowpass the optimized output's lower body for export.

Ported from Mondo-Robotics/PMT TCRS/stair_mppi/mppi_foot_planner_smooth.py @ 88838ca (offline pipeline only; viewer, plum mode, window/trajectory IK, curobo/drake backends and batch CLI not ported). The small joint/body reorder helpers come from the same repo's TCRS/stair_mppi/minimal_mppi_demo.py; the mapping builders were rewritten to resolve joints by name via jnt_qposadr instead of actuator ids.

Attributes

IK_EDGE_HEIGHT_TOL module-attribute

IK_EDGE_HEIGHT_TOL = 0.02

IK_UNSUPPORTED_POINT_SCALE module-attribute

IK_UNSUPPORTED_POINT_SCALE = 0.2

Classes

LowerBodyStabilizerState dataclass

LowerBodyStabilizerState(prev_leg_qpos: ndarray = None, prev_root_pos: ndarray = None, prev_left_contact: bool = False, prev_right_contact: bool = False, left_lock_pos: ndarray = None, left_lock_quat: ndarray = None, right_lock_pos: ndarray = None, right_lock_quat: ndarray = None)

Attributes

prev_leg_qpos class-attribute instance-attribute
prev_leg_qpos: ndarray = None
prev_root_pos class-attribute instance-attribute
prev_root_pos: ndarray = None
prev_left_contact class-attribute instance-attribute
prev_left_contact: bool = False
prev_right_contact class-attribute instance-attribute
prev_right_contact: bool = False
left_lock_pos class-attribute instance-attribute
left_lock_pos: ndarray = None
left_lock_quat class-attribute instance-attribute
left_lock_quat: ndarray = None
right_lock_pos class-attribute instance-attribute
right_lock_pos: ndarray = None
right_lock_quat class-attribute instance-attribute
right_lock_quat: ndarray = None

Methods:

ClipArrays dataclass

ClipArrays(fps: float, dofs_pos: ndarray, links_pos_w: ndarray, links_quat_w: ndarray, links_lin_vel_w: ndarray | None = None)

Input motion clip as plain arrays (orders defined by the adapter's TCRSRobotOptions).

Attributes

fps instance-attribute
fps: float
dofs_pos instance-attribute
dofs_pos: ndarray

(T, J) joint positions in robot_spec.dofs_name order.

links_pos_w: ndarray

(T, K, 3) world link positions in robot_spec.links_name order.

links_quat_w: ndarray

(T, K, 4) world link quaternions, wxyz.

links_lin_vel_w: ndarray | None = None

(T, K, 3) world link linear velocities; finite-differenced when omitted.

Methods:

MotionSet dataclass

MotionSet(dofs_pos: ndarray, links_pos_w: ndarray, links_quat_w: ndarray)

One output motion variant (positions only; velocities are derived downstream).

Attributes

dofs_pos instance-attribute
dofs_pos: ndarray

(T, J) joint positions in robot_spec.dofs_name order.

links_pos_w: ndarray

(T, K, 3) world link positions in robot_spec.links_name order.

links_quat_w: ndarray

(T, K, 4) world link quaternions, wxyz.

Methods:

AdaptResult dataclass

AdaptResult(raw: MotionSet, ghost: MotionSet, optimized: MotionSet, left_contact: ndarray, right_contact: ndarray, ik_pos_errors: ndarray, ik_bad_frame_count: int)

Output of :meth:TCRSMotionAdapter.adapt for one clip placement.

Attributes

raw instance-attribute
raw: MotionSet

SE(2)-placed original motion (no terrain adaptation).

ghost instance-attribute
ghost: MotionSet

Raw joints + terrain-adjusted pelvis z, rebuilt by FK only (no IK).

optimized instance-attribute
optimized: MotionSet

IK-solved motion on terrain; root position, waist and legs zero-phase lowpassed (lowpass_cutoff_hz).

left_contact instance-attribute
left_contact: ndarray

(T,) bool detected left-foot contacts.

right_contact instance-attribute
right_contact: ndarray

(T,) bool detected right-foot contacts.

ik_pos_errors instance-attribute
ik_pos_errors: ndarray

(T,) per-frame max IK position error in millimeters.

ik_bad_frame_count instance-attribute
ik_bad_frame_count: int

Number of frames whose IK error exceeds params.reject_ik_mm.

Methods:

TCRSMotionAdapter

TCRSMotionAdapter(model_path: str, robot_spec: TCRSRobotOptions | None = None, params: TCRSPlannerOptions | None = None)

Terrain-adaptive motion optimizer: adapts flat-ground clips to a terrain.

Loads the MuJoCo model once and reuses the IK solver across :meth:adapt calls; each call allocates fresh MjData so no state leaks between clips.

PARAMETER DESCRIPTION
model_path

Path to the robot's MuJoCo MJCF (floating base + 1-DOF joints).

TYPE: str

robot_spec

Robot geometry/naming; defaults to :class:~eden.options.extensions.motion_optimization.TCRSRobotOptions (the G1 spec).

TYPE: TCRSRobotOptions | None DEFAULT: None

params

Pipeline knobs; defaults to :class:~eden.options.extensions.motion_optimization.TCRSPlannerOptions.

TYPE: TCRSPlannerOptions | None DEFAULT: None

Source code in eden/extensions/motion_optimization/tcrs/pipeline.py
def __init__(
    self,
    model_path: str,
    robot_spec: TCRSRobotOptions | None = None,
    params: TCRSPlannerOptions | None = None,
):
    self.robot_spec = robot_spec if robot_spec is not None else TCRSRobotOptions()
    self.params = params if params is not None else TCRSPlannerOptions()
    self.model = mujoco.MjModel.from_xml_path(str(model_path))
    if self.model.njnt == 0 or self.model.jnt_type[0] != mujoco.mjtJoint.mjJNT_FREE:
        raise ValueError(
            f"The MuJoCo model at '{model_path}' must have a free joint first (floating base); "
            "the pipeline addresses qpos as [root(7), joints...]."
        )
    n_model_joints = self.model.nq - 7
    if len(self.robot_spec.dofs_name) != n_model_joints:
        raise ValueError(
            f"TCRSRobotOptions.dofs_name has {len(self.robot_spec.dofs_name)} joints but the MuJoCo model at "
            f"'{model_path}' has nq-7={n_model_joints} joint DOFs; the clip joint order must cover exactly "
            "the model's non-root joints."
        )
    self.dofs_mapping = build_dofs_mapping(self.model, self.robot_spec.dofs_name)
    self.links_mapping = build_links_mapping(self.model, self.robot_spec.links_name)
    self.ik_solver, self.root_xy_weight, self.root_z_weight = _build_ik_solver(
        self.model, self.robot_spec, log_config=False
    )

Attributes

robot_spec instance-attribute
robot_spec = robot_spec if robot_spec is not None else TCRSRobotOptions()
params instance-attribute
params = params if params is not None else TCRSPlannerOptions()
model instance-attribute
model = mujoco.MjModel.from_xml_path(str(model_path))
dofs_mapping instance-attribute
dofs_mapping = build_dofs_mapping(self.model, self.robot_spec.dofs_name)
links_mapping = build_links_mapping(self.model, self.robot_spec.links_name)

Methods:

adapt
adapt(clip: ClipArrays, terrain: TerrainSource, dx: float = 0.0, dy: float = 0.0, dyaw: float = 0.0) -> AdaptResult

Run the full pipeline for one clip with a given SE(2) placement transform.

PARAMETER DESCRIPTION
clip

Input motion arrays.

TYPE: ClipArrays

terrain

Terrain height source.

TYPE: TerrainSource

dx

X translation applied after the yaw rotation (m).

TYPE: float DEFAULT: 0.0

dy

Y translation applied after the yaw rotation (m).

TYPE: float DEFAULT: 0.0

dyaw

Yaw rotation about the frame-0 pelvis xy (rad).

TYPE: float DEFAULT: 0.0

RETURNS DESCRIPTION
AdaptResult

Raw / ghost / optimized motion sets and IK quality metrics.

Source code in eden/extensions/motion_optimization/tcrs/pipeline.py
2021
2022
2023
2024
2025
2026
2027
2028
2029
2030
2031
2032
2033
2034
2035
2036
2037
2038
2039
2040
2041
2042
2043
2044
2045
2046
2047
2048
2049
2050
2051
2052
2053
2054
2055
2056
2057
2058
2059
2060
2061
2062
2063
2064
2065
2066
2067
2068
2069
2070
2071
2072
2073
2074
2075
2076
2077
2078
2079
2080
2081
2082
2083
2084
2085
2086
2087
2088
2089
2090
2091
2092
2093
2094
2095
2096
2097
2098
2099
2100
2101
2102
2103
2104
2105
2106
2107
2108
2109
2110
2111
2112
2113
2114
2115
2116
2117
2118
2119
2120
2121
2122
2123
2124
2125
2126
2127
2128
2129
2130
2131
2132
2133
2134
2135
2136
2137
2138
2139
2140
2141
2142
2143
2144
2145
2146
2147
2148
2149
2150
2151
2152
2153
2154
2155
2156
2157
2158
2159
2160
2161
2162
2163
2164
2165
2166
2167
2168
2169
2170
2171
2172
2173
2174
2175
2176
2177
2178
2179
2180
2181
2182
2183
2184
2185
2186
2187
2188
2189
2190
2191
2192
2193
2194
2195
2196
2197
2198
2199
2200
2201
2202
2203
2204
2205
2206
2207
2208
2209
2210
2211
2212
2213
2214
2215
2216
2217
2218
2219
2220
2221
2222
2223
2224
2225
2226
2227
2228
2229
2230
2231
2232
2233
2234
2235
2236
2237
2238
2239
2240
2241
2242
2243
2244
2245
2246
2247
2248
2249
2250
2251
2252
2253
2254
2255
2256
2257
2258
2259
2260
2261
2262
2263
2264
2265
2266
2267
2268
2269
2270
2271
2272
2273
2274
2275
2276
2277
2278
2279
2280
2281
2282
2283
2284
2285
2286
2287
2288
2289
2290
2291
2292
2293
2294
2295
2296
2297
2298
2299
2300
2301
2302
2303
2304
2305
2306
2307
2308
2309
2310
2311
2312
2313
2314
2315
2316
2317
2318
2319
2320
2321
2322
2323
2324
2325
2326
2327
2328
2329
2330
2331
2332
2333
2334
2335
2336
2337
2338
2339
2340
2341
2342
2343
2344
2345
2346
2347
2348
2349
2350
2351
2352
2353
2354
2355
2356
2357
2358
2359
2360
2361
2362
2363
2364
2365
2366
2367
2368
2369
2370
2371
2372
2373
2374
2375
2376
2377
2378
2379
2380
2381
2382
2383
2384
2385
2386
2387
2388
2389
2390
2391
2392
2393
def adapt(
    self,
    clip: ClipArrays,
    terrain: TerrainSource,
    dx: float = 0.0,
    dy: float = 0.0,
    dyaw: float = 0.0,
) -> AdaptResult:
    """Run the full pipeline for one clip with a given SE(2) placement transform.

    Parameters
    ----------
    clip : ClipArrays
        Input motion arrays.
    terrain : TerrainSource
        Terrain height source.
    dx : float
        X translation applied after the yaw rotation (m).
    dy : float
        Y translation applied after the yaw rotation (m).
    dyaw : float
        Yaw rotation about the frame-0 pelvis xy (rad).

    Returns
    -------
    AdaptResult
        Raw / ghost / optimized motion sets and IK quality metrics.
    """
    robot_spec = self.robot_spec
    params = self.params

    # ===== Load motion =====
    motion = MotionClip(
        robot_spec=robot_spec,
        fps=clip.fps,
        dofs_pos=clip.dofs_pos,
        links_pos_w=clip.links_pos_w,
        links_quat_w=clip.links_quat_w,
        links_lin_vel_w=clip.links_lin_vel_w,
        contact_z_threshold=params.contact_z_threshold,
        contact_speed_threshold=params.contact_speed_threshold,
    )
    phase_clock_trusted = _should_trust_phase_clock(motion.phase_params)
    if motion.phase_params is not None:
        pp = motion.phase_params
        trust_str = "trusted" if phase_clock_trusted else "disabled"
        en.logger.info(f"  Phase clock: type={pp.gait_type} fit={pp.fit_quality:.2f} [{trust_str}]")

    # ===== Apply SE(2) transform (dx, dy, dyaw) =====
    _apply_se2_transform(motion, dx, dy, dyaw, robot_spec.pelvis_index)

    N = motion.n_frames
    fps = float(motion.fps)

    # ===== Terrain reference =====
    ref_builder = TerrainReference(
        terrain=terrain,
        motion=motion,
        lookahead=params.lookahead,
        smoothing_alpha=params.smoothing_alpha,
        toe_offset_x=robot_spec.toe_probe_offset_x,
        heel_offset_x=params.heel_offset_x,
        mid_offset_x=params.mid_offset_x,
        swing_floor_margin=params.swing_floor_margin,
    )

    # ===== Swing trajectory optimization =====
    midfoot_contact_z = motion.foot_nominal_z + float(np.asarray(robot_spec.mid_foot_offset)[2])
    opt_params = MppiFootParams(
        n_knots=params.mppi_n_knots,
        n_samples=params.mppi_n_samples,
        n_iterations=params.mppi_n_iterations,
        temperature=params.mppi_temperature,
        w_track=params.w_track,
        w_terrain=params.w_terrain,
        h_clearance=params.h_clearance,
        h_contact_z=midfoot_contact_z,
        seed=params.seed,
    )

    pelvis_quats = motion.links_quat[:, robot_spec.pelvis_index]

    # ===== Enforce lateral separation: prevent L/R foot xy crossover =====
    _min_lateral_sep = 0.02
    _pelvis_xy = motion.links_pos[:, robot_spec.pelvis_index, :2]
    for i in range(len(pelvis_quats)):
        yaw = float(yaw_from_quat(pelvis_quats[i]))
        lat = np.array([-np.sin(yaw), np.cos(yaw)])
        lf_lat = np.dot(ref_builder.left_precomputed[i, :2] - _pelvis_xy[i], lat)
        rf_lat = np.dot(ref_builder.right_precomputed[i, :2] - _pelvis_xy[i], lat)
        if lf_lat - rf_lat < _min_lateral_sep:
            deficit = _min_lateral_sep - (lf_lat - rf_lat)
            ref_builder.left_precomputed[i, :2] += 0.5 * deficit * lat
            ref_builder.right_precomputed[i, :2] -= 0.5 * deficit * lat

    mppi_lfoot, mppi_lfoot_quats = precompute_mppi_foot(
        ref_builder.left_precomputed,
        ref_builder.left_precomputed_quats,
        motion.left_contact,
        motion.foot_nominal_z,
        terrain,
        opt_params,
        motion.fps,
        robot_spec=robot_spec,
    )
    mppi_rfoot, mppi_rfoot_quats = precompute_mppi_foot(
        ref_builder.right_precomputed,
        ref_builder.right_precomputed_quats,
        motion.right_contact,
        motion.foot_nominal_z,
        terrain,
        opt_params,
        motion.fps,
        robot_spec=robot_spec,
    )

    # Post-MPPI lateral clamp: MPPI may re-introduce crossover via terrain avoidance,
    # so the same 2 cm separation as before MPPI is applied to its output.
    _min_lateral_sep = 0.02
    _pelvis_xy = motion.links_pos[:, robot_spec.pelvis_index, :2]
    _n_post_fixed = 0
    for i in range(N):
        yaw = float(yaw_from_quat(pelvis_quats[i]))
        lat = np.array([-np.sin(yaw), np.cos(yaw)])
        lf_lat = np.dot(mppi_lfoot[i, :2] - _pelvis_xy[i], lat)
        rf_lat = np.dot(mppi_rfoot[i, :2] - _pelvis_xy[i], lat)
        if lf_lat - rf_lat < _min_lateral_sep:
            deficit = _min_lateral_sep - (lf_lat - rf_lat)
            mppi_lfoot[i, :2] += 0.5 * deficit * lat
            mppi_rfoot[i, :2] -= 0.5 * deficit * lat
            _n_post_fixed += 1
    if _n_post_fixed > 0:
        en.logger.info(f"  Post-MPPI lateral clamp: {_n_post_fixed}/{N} frames")

    # ===== Post-MPPI side-wall collision fix =====
    # Probe horizontally (terrain.foot_side_penetration) and push feet away from terrain side walls.
    if hasattr(terrain, "foot_side_penetration"):
        _n_side_fixed = 0
        for i in range(N):
            yaw = float(yaw_from_quat(pelvis_quats[i]))
            fx, fy = np.cos(yaw), np.sin(yaw)
            # Left foot
            px_l, py_l = _push_side(terrain, mppi_lfoot[i], fx, fy)
            if abs(px_l) > 1e-4 or abs(py_l) > 1e-4:
                mppi_lfoot[i, 0] += px_l
                mppi_lfoot[i, 1] += py_l
                _n_side_fixed += 1
            # Right foot
            px_r, py_r = _push_side(terrain, mppi_rfoot[i], fx, fy)
            if abs(px_r) > 1e-4 or abs(py_r) > 1e-4:
                mppi_rfoot[i, 0] += px_r
                mppi_rfoot[i, 1] += py_r
                _n_side_fixed += 1
        if _n_side_fixed > 0:
            en.logger.info(f"  Post-MPPI side-wall fix: {_n_side_fixed} foot-frames pushed")

    # ===== Pelvis z: construct target externally, filter only smooths =====
    #   support_z   — terrain-driven weighted foot height
    #   z_anchor    — support_z + nominal clearance
    #   z_style     — zero-mean relative pelvis variation from raw motion
    #   z_target    — z_anchor + z_style, clamped by reachability
    root_z_filter = SupportAwareRootZFilter(
        alpha_up=0.5,
        alpha_down=0.35,
        max_delta=0.035,
    )
    raw_pelvis = motion.links_pos[:, robot_spec.pelvis_index].copy()
    nominal_root_clearance = float(motion.pelvis_height_above_foot)
    raw_pelvis_z0 = float(raw_pelvis[0, 2])
    raw_style_z = raw_pelvis[:, 2] - raw_pelvis_z0
    max_reach = robot_spec.max_reach

    pelvis_z_out = np.zeros(N, dtype=np.float64)
    root_z_filter.reset()
    for i in range(N):
        lw, rw, is_flight = _support_state(motion, i, phase_clock_trusted)

        w_sum = lw + rw
        if w_sum > 1e-6:
            wn = np.array([lw, rw]) / w_sum
            support_z = float(wn[0] * mppi_lfoot[i, 2] + wn[1] * mppi_rfoot[i, 2])
        else:
            support_z = float(terrain.height_at(float(raw_pelvis[i, 0]), float(raw_pelvis[i, 1])))

        z_anchor = support_z + nominal_root_clearance
        z_target = z_anchor + float(raw_style_z[i])

        z_upper_candidates = []
        if lw > 1e-3:
            z_upper_candidates.append(_pelvis_reach_z_upper_bound(raw_pelvis[i, :2], mppi_lfoot[i], lw, max_reach))
        if rw > 1e-3:
            z_upper_candidates.append(_pelvis_reach_z_upper_bound(raw_pelvis[i, :2], mppi_rfoot[i], rw, max_reach))
        if not z_upper_candidates:
            z_upper_candidates.extend(
                [
                    _pelvis_reach_z_upper_bound(raw_pelvis[i, :2], mppi_lfoot[i], 0.0, max_reach),
                    _pelvis_reach_z_upper_bound(raw_pelvis[i, :2], mppi_rfoot[i], 0.0, max_reach),
                ]
            )
        z_target = min(z_target, min(z_upper_candidates))

        pelvis_z_out[i] = root_z_filter.step(z_target, is_flight=is_flight)

    pelvis_positions = raw_pelvis.copy()
    pelvis_positions[:, 2] = pelvis_z_out

    # ===== Collision resolution =====
    if params.leg_collision_iters > 0:
        mppi_lfoot, mppi_rfoot = _resolve_lower_leg_foot_collisions(
            motion,
            pelvis_positions,
            mppi_lfoot,
            mppi_rfoot,
            mppi_lfoot_quats,
            mppi_rfoot_quats,
            motion.left_contact,
            motion.right_contact,
            terrain,
            motion.foot_nominal_z,
            n_iters=params.leg_collision_iters,
            shin_radius=params.shin_collision_radius,
            foot_radius=params.foot_collision_radius,
            collision_margin=params.leg_collision_margin,
            push_gain=params.collision_push_gain,
            max_push_per_frame=params.collision_max_push,
            pelvis_side_margin=params.pelvis_side_margin,
            pelvis_side_gain=params.pelvis_side_gain,
        )

    # Remove stance-foot target drift: each stance run keeps the first grounded stance xyz.
    if params.stance_latch:
        mppi_lfoot = _latch_stance_foot_targets(
            mppi_lfoot,
            motion.left_contact,
            terrain,
            motion.foot_nominal_z,
            ground_tol=params.stance_latch_ground_tol,
        )
        mppi_rfoot = _latch_stance_foot_targets(
            mppi_rfoot,
            motion.right_contact,
            terrain,
            motion.foot_nominal_z,
            ground_tol=params.stance_latch_ground_tol,
        )

    left_point_scales = _precompute_contact_point_scales(
        mppi_lfoot,
        mppi_lfoot_quats,
        motion.left_contact,
        terrain,
        robot_spec,
    )
    right_point_scales = _precompute_contact_point_scales(
        mppi_rfoot,
        mppi_rfoot_quats,
        motion.right_contact,
        terrain,
        robot_spec,
    )

    # ===== MuJoCo model + IK setup (model/solver reused from __init__) =====
    model = self.model
    data = mujoco.MjData(model)
    dofs_mapping = self.dofs_mapping
    links_mapping = self.links_mapping
    ik_solver = self.ik_solver
    root_xy_weight = self.root_xy_weight
    root_z_weight = self.root_z_weight

    # Ghost FK model (for offset/ghost recording)
    ghost_model, ghost_data = _build_ghost_fk_model(model)

    n_links_export = len(robot_spec.links_name)
    n_dofs = model.nq - 7  # 29 for G1

    # Buffers: raw is read from motion directly
    # Ghost (offset): raw joints + terrain-adjusted pelvis z
    ghost_dofs_pos = np.zeros((N, n_dofs), dtype=np.float64)
    ghost_links_pos = np.zeros((N, n_links_export, 3), dtype=np.float64)
    ghost_links_quat = np.zeros((N, n_links_export, 4), dtype=np.float64)
    # Optimized (MPPI + IK)
    opt_dofs_pos = np.zeros((N, n_dofs), dtype=np.float64)
    opt_links_pos = np.zeros((N, n_links_export, 3), dtype=np.float64)
    opt_links_quat = np.zeros((N, n_links_export, 4), dtype=np.float64)

    frame_targets, locked_double_stance = _build_frame_ik_targets(
        motion,
        pelvis_positions,
        mppi_lfoot,
        mppi_rfoot,
        mppi_lfoot_quats,
        mppi_rfoot_quats,
        dofs_mapping,
        ik_solver,
        left_point_scales,
        right_point_scales,
        root_xy_weight,
        root_z_weight,
        phase_clock_trusted,
    )

    solved_states, _ = _solve_ik_state_sequence(
        params,
        data,
        ik_solver,
        frame_targets,
        locked_double_stance,
        motion.left_contact[:N],
        motion.right_contact[:N],
        init_progress_prefix=None,
        init_progress_every=200,
        lb_progress_prefix="    [lb-stab]",
        lb_progress_every=200,
    )
    ik_pos_errors = _record_solved_sequence(
        data,
        ghost_model,
        ghost_data,
        ik_solver,
        frame_targets,
        solved_states,
        links_mapping,
        dofs_mapping,
        pelvis_positions,
        opt_dofs_pos,
        opt_links_pos,
        opt_links_quat,
        ghost_dofs_pos,
        ghost_links_pos,
        ghost_links_quat,
        motion,
    )

    # IK error summary
    ik_pos_errors = np.array(ik_pos_errors)
    en.logger.info(
        f"  IK summary: p50={np.percentile(ik_pos_errors, 50):.0f}mm "
        f"p95={np.percentile(ik_pos_errors, 95):.0f}mm "
        f"max={ik_pos_errors.max():.0f}mm  "
        f"bad(>{params.reject_ik_mm:.0f}mm)={np.sum(ik_pos_errors > float(params.reject_ik_mm))}/{N}"
    )

    raw_dofs_pos = motion.dofs_pos[:N].copy()
    raw_links_pos = motion.links_pos[:N].copy()
    raw_links_quat = motion.links_quat[:N].copy()

    # Export lowpass on the optimized output only (raw and ghost stay unfiltered,
    # matching the upstream batch export path).
    effective_cutoff = params.lowpass_cutoff_hz
    if effective_cutoff > 0:
        opt_dofs_pos, opt_links_pos, opt_links_quat = filter_lower_body_motion_for_export(
            model,
            robot_spec,
            dofs_mapping,
            links_mapping,
            opt_dofs_pos,
            opt_links_pos,
            opt_links_quat,
            fps,
            cutoff_hz=effective_cutoff,
            order=params.lowpass_order,
        )

    return AdaptResult(
        raw=MotionSet(dofs_pos=raw_dofs_pos, links_pos_w=raw_links_pos, links_quat_w=raw_links_quat),
        ghost=MotionSet(dofs_pos=ghost_dofs_pos, links_pos_w=ghost_links_pos, links_quat_w=ghost_links_quat),
        optimized=MotionSet(dofs_pos=opt_dofs_pos, links_pos_w=opt_links_pos, links_quat_w=opt_links_quat),
        left_contact=motion.left_contact.copy(),
        right_contact=motion.right_contact.copy(),
        ik_pos_errors=ik_pos_errors.copy(),
        ik_bad_frame_count=int(np.sum(ik_pos_errors > float(params.reject_ik_mm))),
    )

Functions:

build_dofs_mapping

build_dofs_mapping(model: MjModel, dofs_name: list[str]) -> list[int]

Clip joint column i -> index into qpos[7:], resolved by joint name via jnt_qposadr.

Source code in eden/extensions/motion_optimization/tcrs/pipeline.py
def build_dofs_mapping(model: mujoco.MjModel, dofs_name: list[str]) -> list[int]:
    """Clip joint column i -> index into qpos[7:], resolved by joint name via jnt_qposadr."""
    mapping = []
    for name in dofs_name:
        jid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, name)
        if jid == -1:
            raise ValueError(f"Joint '{name}' not found in the MuJoCo model")
        # int() coercion: `x in (enum, ...)` compares with the pybind enum on the left, and
        # mujoco >= 3.12's enums answer False to a numpy scalar there (plain `x == enum` is fine).
        if int(model.jnt_type[jid]) not in (mujoco.mjtJoint.mjJNT_HINGE, mujoco.mjtJoint.mjJNT_SLIDE):
            raise ValueError(f"Joint '{name}' is not a 1-DOF joint")
        mapping.append(int(model.jnt_qposadr[jid]) - 7)
    return mapping
build_links_mapping(model: MjModel, links_name: list[str]) -> list[int]

Build mapping from clip link index to MuJoCo body id: mapping[clip_idx] = mujoco_body_id.

Source code in eden/extensions/motion_optimization/tcrs/pipeline.py
def build_links_mapping(model: mujoco.MjModel, links_name: list[str]) -> list[int]:
    """Build mapping from clip link index to MuJoCo body id: ``mapping[clip_idx] = mujoco_body_id``."""
    mapping = []
    for name in links_name:
        bid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, name)
        if bid == -1:
            raise ValueError(f"Body '{name}' not found in MuJoCo model")
        mapping.append(bid)
    return mapping

reorder_to_mujoco

reorder_to_mujoco(values, mapping)

Reorder a (n_dofs,) clip-order vector into MuJoCo qpos[7:] order.

Source code in eden/extensions/motion_optimization/tcrs/pipeline.py
def reorder_to_mujoco(values, mapping):
    """Reorder a (n_dofs,) clip-order vector into MuJoCo qpos[7:] order."""
    out = np.zeros_like(values)
    for policy_idx, mujoco_idx in enumerate(mapping):
        out[mujoco_idx] = values[policy_idx]
    return out

reorder_from_mujoco

reorder_from_mujoco(values, mapping)

Reorder a (n_dofs,) MuJoCo qpos[7:] vector into clip order.

Source code in eden/extensions/motion_optimization/tcrs/pipeline.py
def reorder_from_mujoco(values, mapping):
    """Reorder a (n_dofs,) MuJoCo qpos[7:] vector into clip order."""
    out = np.zeros_like(values)
    for policy_idx, mujoco_idx in enumerate(mapping):
        out[policy_idx] = values[mujoco_idx]
    return out

update_robot_pose

update_robot_pose(data, root_pos, root_quat, dofs_qpos)

Write a root pose + joint vector into a MuJoCo data (velocities/controls zeroed).

Source code in eden/extensions/motion_optimization/tcrs/pipeline.py
def update_robot_pose(data, root_pos, root_quat, dofs_qpos):
    """Write a root pose + joint vector into a MuJoCo data (velocities/controls zeroed)."""
    data.qpos[:3] = np.asarray(root_pos, dtype=np.float64)
    data.qpos[3:7] = np.asarray(root_quat, dtype=np.float64)
    data.qpos[7:] = np.asarray(dofs_qpos, dtype=np.float64)
    data.qvel[:] = 0.0
    data.ctrl[:] = 0.0

precompute_mppi_foot

precompute_mppi_foot(precomputed_ref, precomputed_quats, contact_mask, foot_nominal_z, terrain, mppi_params, fps, robot_spec)

Run MPPI trajectory optimization over all swing phases of one foot.

Core idea: plan on the mid-foot virtual center (midpoint of toe and heel) rather than the ankle origin. The planning point sits closer to the ground contact surface, reducing toe/heel penetration at step edges. After optimization the mid-foot result is transformed back to ankle coordinates for IK.

A swing between stances at the same height (< 1.5 cm apart) whose reference clears the terrain is kept as it is.

PARAMETER DESCRIPTION
precomputed_ref

(N, 3) ankle reference positions from TerrainReference.

TYPE: ndarray

precomputed_quats

(N, 4) ankle reference quaternions (wxyz) from TerrainReference.

TYPE: ndarray

contact_mask

(N,) boolean contact mask (True = stance, False = swing).

TYPE: ndarray

foot_nominal_z

Nominal ankle z height on flat ground (m).

TYPE: float

terrain

Terrain height source.

TYPE: TerrainSource

mppi_params

Optimizer parameters.

TYPE: MppiFootParams

fps

Motion frame rate (Hz).

TYPE: float

robot_spec

Robot geometry; provides the mid-foot offset for the ankle<->midfoot conversion.

TYPE: TCRSRobotOptions

RETURNS DESCRIPTION
optimized_pos

(N, 3) ankle positions (stance frames unchanged, optimized swing frames overwritten).

TYPE: ndarray

optimized_quats

(N, 4) quaternions: yaw-only during stance (sole stays level), the reference ones during swing, slerp-blended over up to 10 frames at each stance boundary.

TYPE: ndarray

Source code in eden/extensions/motion_optimization/tcrs/pipeline.py
def precompute_mppi_foot(
    precomputed_ref,
    precomputed_quats,
    contact_mask,
    foot_nominal_z,
    terrain,
    mppi_params,
    fps,
    robot_spec,
):
    """Run MPPI trajectory optimization over all swing phases of one foot.

    Core idea: plan on the mid-foot virtual center (midpoint of toe and heel)
    rather than the ankle origin. The planning point sits closer to the ground
    contact surface, reducing toe/heel penetration at step edges. After
    optimization the mid-foot result is transformed back to ankle coordinates
    for IK.

    A swing between stances at the same height (< 1.5 cm apart) whose reference
    clears the terrain is kept as it is.

    Parameters
    ----------
    precomputed_ref : np.ndarray
        (N, 3) ankle reference positions from TerrainReference.
    precomputed_quats : np.ndarray
        (N, 4) ankle reference quaternions (wxyz) from TerrainReference.
    contact_mask : np.ndarray
        (N,) boolean contact mask (True = stance, False = swing).
    foot_nominal_z : float
        Nominal ankle z height on flat ground (m).
    terrain : TerrainSource
        Terrain height source.
    mppi_params : MppiFootParams
        Optimizer parameters.
    fps : float
        Motion frame rate (Hz).
    robot_spec : TCRSRobotOptions
        Robot geometry; provides the mid-foot offset for the ankle<->midfoot conversion.

    Returns
    -------
    optimized_pos : np.ndarray
        (N, 3) ankle positions (stance frames unchanged, optimized swing frames overwritten).
    optimized_quats : np.ndarray
        (N, 4) quaternions: yaw-only during stance (sole stays level), the reference ones
        during swing, slerp-blended over up to 10 frames at each stance boundary.
    """
    opt = MppiFootOptimizer(mppi_params)
    n = len(contact_mask)
    out_pos = precomputed_ref.copy()
    out_quats = precomputed_quats.copy()

    # ===== Step 1: lock stance-phase quaternions to pure yaw =====
    # On discrete stair-like terrain, the stance foot stays level by default (yaw only).
    # This avoids bogus pitch/roll from finite-difference normal estimates at step edges.
    stance_phases = _find_stance_phases(contact_mask)
    for st, en_ in stance_phases:
        yaw = float(yaw_from_quat(precomputed_quats[st]))
        locked_quat = xyz_to_quat(np.array([0.0, 0.0, yaw]))  # yaw-only: sole stays level
        for i in range(st, en_):
            out_quats[i] = locked_quat

    # ===== Step 2: run MPPI optimization for each swing phase =====
    phases = _find_swing_phases(contact_mask)
    _n_skipped = 0

    if len(phases) > 0:
        en.logger.info(f"    MPPI optimizing {len(phases)} swing phases...")

    for phase_idx, (swing_start, swing_end) in enumerate(phases):
        if phase_idx == 0 or (phase_idx + 1) % 5 == 0 or phase_idx == len(phases) - 1:
            en.logger.debug(f"      phase {phase_idx + 1}/{len(phases)} frames=[{swing_start},{swing_end})")
        liftoff_idx = swing_start - 1  # liftoff frame = the frame before the swing starts
        landing_idx = swing_end  # landing frame = the first frame after the swing ends

        # Bounds check: skip optimization if the liftoff/landing frame is out of range.
        # Swing frames keep the original mocap quaternions (precomputed_quats), not overwritten.
        if liftoff_idx < 0 or landing_idx >= n:
            continue

        # ===== Conditional skip: preserve original motion when possible =====
        # Check height difference between consecutive stances
        liftoff_z = float(out_pos[liftoff_idx, 2])
        landing_z = float(out_pos[landing_idx, 2])
        height_diff = abs(landing_z - liftoff_z)

        if height_diff < 0.015:  # Same level (< 1.5cm height change)
            # Check: does the original swing trajectory collide with terrain?
            _has_collision = False
            _clearance_margin = mppi_params.h_clearance if hasattr(mppi_params, "h_clearance") else 0.02
            for _k in range(swing_start, swing_end):
                _fx, _fy, _fz = precomputed_ref[_k]
                _tz = float(terrain.height_at(_fx, _fy))
                if _fz < _tz + foot_nominal_z + _clearance_margin:
                    _has_collision = True
                    break

            if not _has_collision:
                # SKIP: no height change, no collision → keep original trajectory
                _n_skipped += 1
                continue

        n_swing = swing_end - swing_start  # number of frames in the swing phase

        # --- Swing quaternions: keep the original mocap quaternions, tracking the
        # reference motion's pitch/roll ---
        # The stance phase was locked to pure yaw (level sole) in step 1;
        # the swing phase keeps the original mocap values from precomputed_quats.
        # Slerp over blend_frames frames at the start and end of the swing to avoid
        # quaternion jumps at the stance (pure yaw) -> swing (mocap) boundary.
        blend_frames = min(10, n_swing // 2)  # number of transition frames, at most half a swing
        if blend_frames > 0:
            q_liftoff = out_quats[liftoff_idx]  # pure yaw (locked during stance)
            q_landing = out_quats[landing_idx]  # pure yaw (locked during stance)
            # Swing start: transition from liftoff pure yaw to mocap
            for b in range(blend_frames):
                alpha = (b + 1) / (blend_frames + 1)  # 0 -> 1 excluding the endpoints
                idx_b = swing_start + b
                q_mocap = precomputed_quats[idx_b]
                q_blend = slerp_quaternion(q_liftoff, q_mocap, np.array([alpha]))[0]
                out_quats[idx_b] = q_blend
            # Swing end: transition from mocap to landing pure yaw
            for b in range(blend_frames):
                alpha = (b + 1) / (blend_frames + 1)
                idx_b = swing_end - blend_frames + b
                q_mocap = precomputed_quats[idx_b]
                q_blend = slerp_quaternion(q_mocap, q_landing, np.array([alpha]))[0]
                out_quats[idx_b] = q_blend

        # --- Convert ankle coordinates to mid-foot coordinates for MPPI planning ---
        liftoff_midfoot = _ankle_to_midfoot(out_pos[liftoff_idx], out_quats[liftoff_idx], robot_spec)
        landing_midfoot = _ankle_to_midfoot(out_pos[landing_idx], out_quats[landing_idx], robot_spec)

        # Special case: with fewer than 2 swing frames MPPI cannot optimize
        # (at least 2 frames are needed to form a valid trajectory) — use linear
        # interpolation instead.
        if n_swing < 2:
            t_frac = np.linspace(0, 1, n_swing + 2)[1:-1]  # drop the two endpoint samples
            for k, frac in enumerate(t_frac):
                midfoot_interp = (1 - frac) * liftoff_midfoot + frac * landing_midfoot
                out_pos[swing_start + k] = _midfoot_to_ankle(midfoot_interp, out_quats[swing_start + k], robot_spec)
            continue

        swing_duration = n_swing / fps  # swing duration (seconds)
        ref_indices = np.arange(swing_start, swing_end)  # indices of the swing frames

        # --- Build the MPPI tracking target: convert reference ankle positions to
        # mid-foot coordinates ---
        ref_traj_midfoot = np.zeros((n_swing, 3), dtype=np.float64)
        for k, idx in enumerate(ref_indices):
            ref_traj_midfoot[k] = _ankle_to_midfoot(precomputed_ref[idx], out_quats[idx], robot_spec)
        ref_times = np.linspace(0, swing_duration, n_swing)  # time sample points

        # optimize_swing samples xyz jointly: tracking + smoothness costs, plus
        # terrain-penetration and height-ceiling penalties on z.
        try:
            traj_x, traj_y, traj_z = opt.optimize_swing(
                liftoff_midfoot,
                landing_midfoot,
                swing_duration,
                ref_traj_midfoot,
                ref_times,
                terrain,
            )
        except Exception as e:
            en.logger.warning(f"  [MPPI] Warning: swing [{swing_start}:{swing_end}] failed: {e}")
            continue  # keep the original TerrainReference trajectory when optimization fails

        # --- Evaluate the optimized mid-foot trajectory and convert back to ankle
        # coordinates ---
        t_eval = np.linspace(0, swing_duration, n_swing)
        midfoot_x = traj_x.eval(t_eval)
        midfoot_y = traj_y.eval(t_eval)
        midfoot_z = traj_z.eval(t_eval)

        for k, idx in enumerate(ref_indices):
            midfoot_pos = np.array([midfoot_x[k], midfoot_y[k], midfoot_z[k]])
            out_pos[idx] = _midfoot_to_ankle(midfoot_pos, out_quats[idx], robot_spec)  # mid-foot -> ankle

        # --- Blend back toward original ref if landing at same height ---
        # When MPPI was needed (collision on flat ground), blend the end of
        # the swing back toward the precomputed_ref to restore original motion shape.
        if height_diff < 0.015 and n_swing > 6:
            _blend_start = swing_start + int(0.7 * n_swing)
            for _k in range(_blend_start, swing_end):
                _alpha = (_k - _blend_start) / max(swing_end - _blend_start, 1)
                out_pos[_k] = (1.0 - _alpha) * out_pos[_k] + _alpha * precomputed_ref[_k]

    if _n_skipped > 0:
        en.logger.info(f"    Skipped {_n_skipped}/{len(phases)} swing phases (no height change, no collision)")

    return out_pos, out_quats

apply_zero_phase_lowpass

apply_zero_phase_lowpass(traj, fps, cutoff_hz=10.0, order=4)

Zero-phase Butterworth lowpass on a (N,D) or (N,) array.

Uses scipy.signal.filtfilt for zero phase shift. Gracefully handles short sequences: if the signal is too short for the requested filter order, the order is reduced until filtfilt can run, or the original is returned unfiltered.

Source code in eden/extensions/motion_optimization/tcrs/pipeline.py
def apply_zero_phase_lowpass(traj, fps, cutoff_hz=10.0, order=4):
    """Zero-phase Butterworth lowpass on a (N,D) or (N,) array.

    Uses scipy.signal.filtfilt for zero phase shift.
    Gracefully handles short sequences: if the signal is too short for
    the requested filter order, the order is reduced until filtfilt can
    run, or the original is returned unfiltered.
    """
    from scipy.signal import butter, filtfilt

    if cutoff_hz <= 0:
        return traj.copy()
    nyquist = fps / 2.0
    if cutoff_hz >= nyquist:
        return traj.copy()

    traj = np.asarray(traj, dtype=np.float64)
    N = traj.shape[0]

    # filtfilt requires padlen <= N - 1; default padlen = 3 * max(len(a), len(b))
    # For butter order k: len(b) = len(a) = k + 1, so padlen = 3 * (k + 1).
    # If N is too short, reduce order until it fits, or skip.
    eff_order = int(order)
    while eff_order >= 1:
        padlen = 3 * (eff_order + 1)
        if padlen < N:
            break
        eff_order -= 1
    if eff_order < 1:
        return traj.copy()  # too short to filter at all

    b, a = butter(eff_order, cutoff_hz / nyquist, btype="low")

    def _filt1d(x):
        return filtfilt(b, a, x)

    if traj.ndim == 1:
        return _filt1d(traj)
    out = np.empty_like(traj)
    if traj.ndim == 2:
        for j in range(traj.shape[1]):
            out[:, j] = _filt1d(traj[:, j])
    elif traj.ndim == 3:
        for b_idx in range(traj.shape[1]):
            for d in range(traj.shape[2]):
                out[:, b_idx, d] = _filt1d(traj[:, b_idx, d])
    else:
        out = traj.copy()
    return out

filter_motion_for_export

filter_motion_for_export(dofs_pos, links_pos, links_quat, fps, cutoff_hz=10.0, order=4)

Apply zero-phase lowpass to all exported channels consistently.

  • dofs_pos, links_pos: Butterworth filtfilt.
  • links_quat: log-map → filtfilt → exp-map (proper SO(3) smoothing).

(Upstream also recomputed velocities from the filtered channels; the port returns positions/quaternions only.)

Source code in eden/extensions/motion_optimization/tcrs/pipeline.py
def filter_motion_for_export(
    dofs_pos,
    links_pos,
    links_quat,
    fps,
    cutoff_hz=10.0,
    order=4,
):
    """Apply zero-phase lowpass to all exported channels consistently.

    - dofs_pos, links_pos: Butterworth filtfilt.
    - links_quat: log-map → filtfilt → exp-map (proper SO(3) smoothing).

    (Upstream also recomputed velocities from the filtered channels; the port
    returns positions/quaternions only.)
    """
    if cutoff_hz <= 0:
        return dofs_pos.copy(), links_pos.copy(), links_quat.copy()

    filt_dofs_pos = apply_zero_phase_lowpass(dofs_pos, fps, cutoff_hz, order)
    filt_links_pos = apply_zero_phase_lowpass(links_pos, fps, cutoff_hz, order)
    filt_links_quat = _smooth_quaternions_logmap(links_quat, fps, cutoff_hz, order)
    return filt_dofs_pos, filt_links_pos, filt_links_quat

filter_lower_body_motion_for_export

filter_lower_body_motion_for_export(model, robot_spec, dofs_mapping, links_mapping, dofs_pos, links_pos, links_quat, fps, cutoff_hz=5.0, order=4)

Smooth root position + waist + leg joints, then rebuild full-body poses by FK.

This preserves shoulder/elbow/wrist motion while making the support chain smoother and keeping exported body poses consistent with the exported joint commands.

Root orientation is intentionally left unchanged. Smoothing pelvis quaternions over long clips can create catastrophic whole-body heading artefacts even when IK position error stays small, because the batch error metric does not constrain root orientation.

(Upstream loaded its own MjModel from the XML and recomputed velocities; the port reuses the caller's model/mappings and returns positions only.)

Source code in eden/extensions/motion_optimization/tcrs/pipeline.py
def filter_lower_body_motion_for_export(
    model,
    robot_spec,
    dofs_mapping,
    links_mapping,
    dofs_pos,
    links_pos,
    links_quat,
    fps,
    cutoff_hz=5.0,
    order=4,
):
    """Smooth root position + waist + leg joints, then rebuild full-body poses by FK.

    This preserves shoulder/elbow/wrist motion while making the support chain
    smoother and keeping exported body poses consistent with the exported
    joint commands.

    Root orientation is intentionally left unchanged. Smoothing pelvis
    quaternions over long clips can create catastrophic whole-body heading
    artefacts even when IK position error stays small, because the batch
    error metric does not constrain root orientation.

    (Upstream loaded its own MjModel from the XML and recomputed velocities;
    the port reuses the caller's model/mappings and returns positions only.)
    """
    dofs_pos = np.asarray(dofs_pos, dtype=np.float64)
    links_pos = np.asarray(links_pos, dtype=np.float64)
    links_quat = np.asarray(links_quat, dtype=np.float64)
    if cutoff_hz <= 0:
        return dofs_pos.copy(), links_pos.copy(), links_quat.copy()

    smooth_names = (
        list(robot_spec.left_leg_dofs_name) + list(robot_spec.right_leg_dofs_name) + list(robot_spec.waist_dofs_name)
    )
    smooth_policy_idx = np.array(
        [robot_spec.dofs_name.index(name) for name in smooth_names],
        dtype=np.int64,
    )

    filt_dofs_pos = dofs_pos.copy()
    filt_dofs_pos[:, smooth_policy_idx] = apply_zero_phase_lowpass(
        dofs_pos[:, smooth_policy_idx],
        fps,
        cutoff_hz,
        order,
    )

    filt_root_pos = apply_zero_phase_lowpass(
        links_pos[:, robot_spec.pelvis_index, :],
        fps,
        cutoff_hz,
        order,
    )
    filt_root_quat = links_quat[:, robot_spec.pelvis_index].copy()

    data = mujoco.MjData(model)

    filt_links_pos = np.zeros_like(links_pos, dtype=np.float64)
    filt_links_quat = np.zeros_like(links_quat, dtype=np.float64)
    for t in range(dofs_pos.shape[0]):
        dofs_qpos = reorder_to_mujoco(filt_dofs_pos[t], dofs_mapping)
        update_robot_pose(
            data,
            root_pos=filt_root_pos[t],
            root_quat=filt_root_quat[t],
            dofs_qpos=dofs_qpos,
        )
        mujoco.mj_forward(model, data)
        for npz_i, mj_bid in enumerate(links_mapping):
            filt_links_pos[t, npz_i] = data.xpos[mj_bid]
            filt_links_quat[t, npz_i] = data.xquat[mj_bid]

    return filt_dofs_pos, filt_links_pos, filt_links_quat

Motion clips and terrain references

Motion clip container, footstep resolution, and terrain-adapted reference trajectories.

:class:MotionClip wraps a mocap segment (joint positions plus world-frame body poses) and derives per-foot contact masks — an adaptive percentile z-threshold combined with a speed gate, stabilized by :func:stabilize_contact_mask — and an optional gait phase clock for gaits with flight phases. :class:FootstepResolver snaps each touchdown to a safe landing pose on the terrain and locks it across the stance phase. :class:TerrainReference builds the terrain-adapted pelvis and foot reference trajectories consumed by the planner (stance-locked footholds, smooth swing interpolation, floor clamping, lateral crossover and step-order enforcement).

Ported from Mondo-Robotics/PMT TCRS/stair_mppi/minimal_mppi_demo.py @ 88838ca (viewer demo, TinyMPPI and MuJoCo helpers not ported).

Classes

MotionClip

MotionClip(robot_spec: TCRSRobotOptions, fps: float, dofs_pos: ndarray, links_pos_w: ndarray, links_quat_w: ndarray, links_lin_vel_w: ndarray | None = None, contact_z_threshold: float = 0.06, contact_speed_threshold: float = 0.5)

A mocap motion segment with derived foot-contact masks and gait phase.

PARAMETER DESCRIPTION
robot_spec

Robot geometry/naming; resolves the pelvis and foot columns of the body arrays and provides the nominal ankle-to-sole height.

TYPE: TCRSRobotOptions

fps

Frame rate of the clip (Hz).

TYPE: float

dofs_pos

(T, n_dofs) joint positions in robot_spec.dofs_name order.

TYPE: ndarray

links_pos_w

(T, n_links, 3) world body positions in robot_spec.links_name order.

TYPE: ndarray

links_quat_w

(T, n_links, 4) world body quaternions [w, x, y, z].

TYPE: ndarray

links_lin_vel_w

(T, n_links, 3) world body linear velocities. Finite-differenced from links_pos_w when omitted.

TYPE: ndarray or None DEFAULT: None

contact_z_threshold

Minimum foot-z threshold for contact detection (m); the adaptive percentile threshold never goes below it.

TYPE: float DEFAULT: 0.06

contact_speed_threshold

Foot speed threshold for contact detection (m/s).

TYPE: float DEFAULT: 0.5

Source code in eden/extensions/motion_optimization/tcrs/motion.py
def __init__(
    self,
    robot_spec: TCRSRobotOptions,
    fps: float,
    dofs_pos: np.ndarray,
    links_pos_w: np.ndarray,
    links_quat_w: np.ndarray,
    links_lin_vel_w: np.ndarray | None = None,
    contact_z_threshold: float = 0.06,
    contact_speed_threshold: float = 0.5,
):
    dofs_pos = np.asarray(dofs_pos)
    links_pos_w = np.asarray(links_pos_w)
    links_quat_w = np.asarray(links_quat_w)
    if fps <= 0.0:
        raise ValueError(f"fps must be > 0, got {fps}")
    if dofs_pos.ndim != 2:
        raise ValueError(f"dofs_pos must be (T, n_dofs), got shape {dofs_pos.shape}")
    n_frames = int(dofs_pos.shape[0])
    if n_frames == 0:
        raise ValueError("Motion clip is empty (0 frames)")
    n_dofs = len(robot_spec.dofs_name)
    n_links = len(robot_spec.links_name)
    if dofs_pos.shape[1] != n_dofs:
        raise ValueError(f"dofs_pos has {dofs_pos.shape[1]} joints but robot_spec.dofs_name has {n_dofs}")
    if links_pos_w.shape != (n_frames, n_links, 3):
        raise ValueError(
            f"links_pos_w must have shape ({n_frames}, {n_links}, 3) to match dofs_pos frames and "
            f"robot_spec.links_name, got {links_pos_w.shape}"
        )
    if links_quat_w.shape != (n_frames, n_links, 4):
        raise ValueError(f"links_quat_w must have shape ({n_frames}, {n_links}, 4), got {links_quat_w.shape}")
    if links_lin_vel_w is not None:
        links_lin_vel_w = np.asarray(links_lin_vel_w)
        if links_lin_vel_w.shape != (n_frames, n_links, 3):
            raise ValueError(
                f"links_lin_vel_w must have shape ({n_frames}, {n_links}, 3), got {links_lin_vel_w.shape}"
            )

    self.robot_spec = robot_spec
    self.fps = float(fps)
    self.n_frames = n_frames
    self.dofs_pos = dofs_pos.copy()
    self.links_pos = links_pos_w.copy()
    self.links_quat = links_quat_w.copy()
    if links_lin_vel_w is not None:
        self.links_lin_vel = links_lin_vel_w.copy()
    else:
        self.links_lin_vel = np.gradient(self.links_pos, axis=0) * self.fps

    pelvis_idx = robot_spec.pelvis_index
    lfoot_idx = robot_spec.left_foot_index
    rfoot_idx = robot_spec.right_foot_index
    lfoot_speed = np.linalg.norm(self.links_lin_vel[:, lfoot_idx], axis=1)
    rfoot_speed = np.linalg.norm(self.links_lin_vel[:, rfoot_idx], axis=1)
    lfoot_z = self.links_pos[:, lfoot_idx, 2]
    rfoot_z = self.links_pos[:, rfoot_idx, 2]

    # Adaptive contact threshold: percentile-based.
    # Use the 25th percentile of foot z + margin as threshold.
    # This is robust to motions where the ground plane is not at z=0.
    all_foot_z = np.concatenate([lfoot_z, rfoot_z])
    _p25 = float(np.percentile(all_foot_z, 25.0))
    _adaptive_z_thresh = max(float(contact_z_threshold), _p25 + 0.02)
    left_contact_raw = (lfoot_z < _adaptive_z_thresh) & (lfoot_speed < float(contact_speed_threshold))
    right_contact_raw = (rfoot_z < _adaptive_z_thresh) & (rfoot_speed < float(contact_speed_threshold))
    self.left_contact = stabilize_contact_mask(left_contact_raw)
    self.right_contact = stabilize_contact_mask(right_contact_raw)

    # Phase clock fitting — for gaits with flight phases (running/hopping)
    self.phase_params = fit_phase_clock(
        lfoot_z,
        rfoot_z,
        lfoot_speed,
        rfoot_speed,
        self.fps,
        contact_z_threshold,
        contact_speed_threshold,
    )
    if self.phase_params is not None:
        self.per_frame_phase = compute_per_frame_phase(self.phase_params, self.n_frames)
        # Do NOT override contact masks — phase clock can misdetect boundaries.
    else:
        self.per_frame_phase = None

    # The robot's known ankle-to-sole height (upstream: 0.06 m, measured
    # from the G1 ankle_roll_link to the sole).
    self.foot_nominal_z = robot_spec.foot_nominal_z

    # Compute nominal pelvis height above foot contact surface (flat ground).
    # This is the target for pelvis height regulation on terrain.
    pelvis_z = self.links_pos[:, pelvis_idx, 2]
    self.pelvis_height_above_foot = float(np.median(pelvis_z) - self.foot_nominal_z)

Attributes

robot_spec instance-attribute
robot_spec = robot_spec
fps instance-attribute
fps = float(fps)
n_frames instance-attribute
n_frames = n_frames
dofs_pos instance-attribute
dofs_pos = dofs_pos.copy()
links_pos = links_pos_w.copy()
links_quat = links_quat_w.copy()
links_lin_vel = links_lin_vel_w.copy()
left_contact instance-attribute
left_contact = stabilize_contact_mask(left_contact_raw)
right_contact instance-attribute
right_contact = stabilize_contact_mask(right_contact_raw)
phase_params instance-attribute
phase_params = fit_phase_clock(lfoot_z, rfoot_z, lfoot_speed, rfoot_speed, self.fps, contact_z_threshold, contact_speed_threshold)
per_frame_phase instance-attribute
per_frame_phase = compute_per_frame_phase(self.phase_params, self.n_frames)
foot_nominal_z instance-attribute
foot_nominal_z = robot_spec.foot_nominal_z
pelvis_height_above_foot instance-attribute
pelvis_height_above_foot = float(np.median(pelvis_z) - self.foot_nominal_z)

Methods:

get_pelvis_pos
get_pelvis_pos(frame: int) -> ndarray
Source code in eden/extensions/motion_optimization/tcrs/motion.py
def get_pelvis_pos(self, frame: int) -> np.ndarray:
    return self.links_pos[frame, self.robot_spec.pelvis_index].copy()
get_future_indices
get_future_indices(frame: int, horizon: int) -> ndarray
Source code in eden/extensions/motion_optimization/tcrs/motion.py
def get_future_indices(self, frame: int, horizon: int) -> np.ndarray:
    return self._future_indices(frame, horizon)
get_future_segment
get_future_segment(frame: int, horizon: int) -> tuple[ndarray, ndarray, ndarray]
Source code in eden/extensions/motion_optimization/tcrs/motion.py
def get_future_segment(self, frame: int, horizon: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    idx = self._future_indices(frame, horizon)
    return (
        self.links_pos[idx, self.robot_spec.pelvis_index].copy(),
        self.links_pos[idx, self.robot_spec.left_foot_index].copy(),
        self.links_pos[idx, self.robot_spec.right_foot_index].copy(),
    )
get_future_contacts
get_future_contacts(frame: int, horizon: int) -> tuple[ndarray, ndarray]
Source code in eden/extensions/motion_optimization/tcrs/motion.py
def get_future_contacts(self, frame: int, horizon: int) -> tuple[np.ndarray, np.ndarray]:
    idx = self._future_indices(frame, horizon)
    return (
        self.left_contact[idx].copy(),
        self.right_contact[idx].copy(),
    )
get_future_support_weights
get_future_support_weights(frame: int, horizon: int) -> tuple[ndarray, ndarray]

Return continuous support weights for pelvis z computation.

When phase clock is available, returns smooth cosine-decay weights. Otherwise, falls back to binary weights from contact masks (1.0/0.0).

Source code in eden/extensions/motion_optimization/tcrs/motion.py
def get_future_support_weights(self, frame: int, horizon: int) -> tuple[np.ndarray, np.ndarray]:
    """Return continuous support weights for pelvis z computation.

    When phase clock is available, returns smooth cosine-decay weights.
    Otherwise, falls back to binary weights from contact masks (1.0/0.0).
    """
    idx = self._future_indices(frame, horizon)
    if self.per_frame_phase is not None:
        return (
            self.per_frame_phase.support_weight_left[idx].copy(),
            self.per_frame_phase.support_weight_right[idx].copy(),
        )
    return (
        self.left_contact[idx].astype(np.float64),
        self.right_contact[idx].astype(np.float64),
    )
get_future_foot_quats
get_future_foot_quats(frame: int, horizon: int) -> tuple[ndarray, ndarray]
Source code in eden/extensions/motion_optimization/tcrs/motion.py
def get_future_foot_quats(self, frame: int, horizon: int) -> tuple[np.ndarray, np.ndarray]:
    idx = self._future_indices(frame, horizon)
    return (
        self.links_quat[idx, self.robot_spec.left_foot_index].copy(),
        self.links_quat[idx, self.robot_spec.right_foot_index].copy(),
    )

FootstepResolver

FootstepResolver(robot_spec: TCRSRobotOptions, terrain: TerrainSource, raw_foot: ndarray, contact_mask: ndarray, foot_nominal_z: float, foot_quats: ndarray | None = None, foot_point_offsets: ndarray | None = None, toe_offset_x: float | None = None, max_landing_shift: float = 0.05)

Footstep resolver: snaps raw mocap ankle positions onto safe terrain support surfaces.

Core responsibilities:

  1. For each touchdown frame, call :meth:_resolve_landing_pose to find the best support placement, constraining the ankle so all foot key points (heel/midfoot/toe) land on the same support surface.
  2. Broadcast the resolved landing pose across the whole stance phase (contact_support) so the foot does not slide while in contact.
PARAMETER DESCRIPTION
robot_spec

Robot geometry; provides the default toe probe offset.

TYPE: TCRSRobotOptions

terrain

Terrain object; must provide height_at(x, y) and height_batch(x, y).

TYPE: TerrainSource

raw_foot

(N, 3) raw mocap ankle positions (world frame).

TYPE: ndarray

contact_mask

(N,) boolean contact mask; True = stance.

TYPE: ndarray

foot_nominal_z

Nominal ankle z height when standing on flat ground (z=0).

TYPE: float

foot_quats

(N, 4) ankle quaternions [w, x, y, z], used to rotate foot offsets into the world frame. Defaults to identity (facing forward).

TYPE: ndarray or None DEFAULT: None

foot_point_offsets

(M,) local x offsets of foot key points relative to the ankle. Each offset is rotated by the quaternion into world coordinates and checked against step edges. Defaults to [toe_offset_x].

TYPE: ndarray or None DEFAULT: None

toe_offset_x

Toe x offset relative to the ankle (local frame). Defaults to robot_spec.toe_probe_offset_x.

TYPE: float or None DEFAULT: None

max_landing_shift

Maximum search displacement along the foot forward direction in :meth:_resolve_landing_pose (m).

TYPE: float DEFAULT: 0.05

Source code in eden/extensions/motion_optimization/tcrs/motion.py
def __init__(
    self,
    robot_spec: TCRSRobotOptions,
    terrain: TerrainSource,
    raw_foot: np.ndarray,
    contact_mask: np.ndarray,
    foot_nominal_z: float,
    foot_quats: np.ndarray | None = None,
    foot_point_offsets: np.ndarray | None = None,
    toe_offset_x: float | None = None,
    max_landing_shift: float = 0.05,
):
    if toe_offset_x is None:
        toe_offset_x = robot_spec.toe_probe_offset_x
    self.robot_spec = robot_spec
    self.terrain = terrain
    self.raw_foot = np.asarray(raw_foot, dtype=np.float64)
    self.contact_mask = np.asarray(contact_mask, dtype=bool)
    self.foot_nominal_z = float(foot_nominal_z)
    self.toe_offset_x = max(float(toe_offset_x), 0.0)
    self.max_landing_shift = max(float(max_landing_shift), 0.005)
    self.n_frames = int(self.contact_mask.shape[0])
    if foot_quats is not None:
        self.foot_quats = np.asarray(foot_quats, dtype=np.float64)
    else:
        # Without quaternions, default to facing forward (identity quaternion).
        self.foot_quats = np.tile(
            np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float64),
            (self.n_frames, 1),
        )
    if foot_point_offsets is not None:
        self.foot_point_offsets = np.asarray(foot_point_offsets, dtype=np.float64)
    else:
        self.foot_point_offsets = np.array([toe_offset_x], dtype=np.float64)

    # Detect touchdown events: rising edges from non-contact to contact.
    # contact_mask[i] = True and contact_mask[i-1] = False → frame i is a touchdown frame.
    self.touchdown = self.contact_mask & ~np.r_[False, self.contact_mask[:-1]]
    # If contact frames exist but no rising edge was found (the whole segment is
    # in contact), mark the first contact frame as a touchdown.
    if np.any(self.contact_mask) and not np.any(self.touchdown):
        first_contact = int(np.flatnonzero(self.contact_mask)[0])
        self.touchdown[first_contact] = True

    # For each touchdown frame, resolve a safe landing position
    # (x constrained to the step's safe interval, z = terrain height + foot_nominal_z).
    self.touchdown_support = np.full((self.n_frames, 3), np.nan, dtype=np.float64)
    for td in np.flatnonzero(self.touchdown):
        self.touchdown_support[td] = self._resolve_landing_pose(self.raw_foot[td], self.foot_quats[td])

    # Build prev/next contact-frame index maps (used during swing to find adjacent stance frames).
    (
        self.prev_contact_idx,  # prev_contact_idx[i] = nearest contact frame index before frame i
        self.next_contact_idx,  # next_contact_idx[i] = nearest contact frame index after frame i
    ) = self._build_prev_next_indices(self.contact_mask)
    # Build per-frame "next touchdown frame" indices (wrap=True: wraps to the start past the end).
    self.next_touchdown_idx = self._build_next_event_indices(self.touchdown, wrap=True)

    # Broadcast touchdown landing poses across the whole stance phase:
    # contact_support[i] = the locked position of that stance segment.
    self.contact_support = np.full((self.n_frames, 3), np.nan, dtype=np.float64)
    self._build_contact_support()

Attributes

robot_spec instance-attribute
robot_spec = robot_spec
terrain instance-attribute
terrain = terrain
raw_foot instance-attribute
raw_foot = np.asarray(raw_foot, dtype=np.float64)
contact_mask instance-attribute
contact_mask = np.asarray(contact_mask, dtype=bool)
foot_nominal_z instance-attribute
foot_nominal_z = float(foot_nominal_z)
toe_offset_x instance-attribute
toe_offset_x = max(float(toe_offset_x), 0.0)
max_landing_shift instance-attribute
max_landing_shift = max(float(max_landing_shift), 0.005)
n_frames instance-attribute
n_frames = int(self.contact_mask.shape[0])
foot_quats instance-attribute
foot_quats = np.asarray(foot_quats, dtype=np.float64)
foot_point_offsets instance-attribute
foot_point_offsets = np.asarray(foot_point_offsets, dtype=np.float64)
touchdown instance-attribute
touchdown = self.contact_mask & ~np.r_[False, self.contact_mask[:-1]]
touchdown_support instance-attribute
touchdown_support = np.full((self.n_frames, 3), np.nan, dtype=np.float64)
next_touchdown_idx instance-attribute
next_touchdown_idx = self._build_next_event_indices(self.touchdown, wrap=True)
contact_support instance-attribute
contact_support = np.full((self.n_frames, 3), np.nan, dtype=np.float64)

Methods:

prev_support_pose
prev_support_pose(frame_idx: int) -> ndarray | None

Query the nearest support position before frame frame_idx.

Used during swing to determine the liftoff position: the contact_support of the last contact frame before the swing is the liftoff point.

RETURNS DESCRIPTION
ndarray or None

(3,) position array [x, y, z], or None (no preceding support frame).

Source code in eden/extensions/motion_optimization/tcrs/motion.py
def prev_support_pose(self, frame_idx: int) -> np.ndarray | None:
    """Query the nearest support position before frame ``frame_idx``.

    Used during swing to determine the liftoff position: the contact_support
    of the last contact frame before the swing is the liftoff point.

    Returns
    -------
    np.ndarray or None
        (3,) position array [x, y, z], or None (no preceding support frame).
    """
    idx = int(frame_idx) % self.n_frames
    prev_contact = int(self.prev_contact_idx[idx])
    if prev_contact < 0:
        return None
    pose = self.contact_support[prev_contact]
    if not np.isfinite(pose[2]):
        return None
    return pose.copy()
next_touchdown_pose
next_touchdown_pose(frame_idx: int) -> ndarray | None

Query the nearest touchdown (landing) position after frame frame_idx.

Used during swing to determine the landing target: the safe position the foot must reach at the end of the swing.

RETURNS DESCRIPTION
ndarray or None

(3,) position array [x, y, z], or None (no subsequent touchdown).

Source code in eden/extensions/motion_optimization/tcrs/motion.py
def next_touchdown_pose(self, frame_idx: int) -> np.ndarray | None:
    """Query the nearest touchdown (landing) position after frame ``frame_idx``.

    Used during swing to determine the landing target: the safe position the
    foot must reach at the end of the swing.

    Returns
    -------
    np.ndarray or None
        (3,) position array [x, y, z], or None (no subsequent touchdown).
    """
    idx = int(frame_idx) % self.n_frames
    td = int(self.next_touchdown_idx[idx])
    if td < 0:
        return None
    pose = self.touchdown_support[td]
    if not np.isfinite(pose[2]):
        return None
    return pose.copy()

TerrainReference

TerrainReference(terrain: TerrainSource, motion: MotionClip, lookahead: float = 0.0, smoothing_alpha: float = 0.08, toe_offset_x: float | None = None, heel_offset_x: float = -0.04, mid_offset_x: float = 0.1, swing_floor_margin: float = 0.003, soft_reach_limit: float = 0.7, max_landing_shift: float = 0.05)

Terrain reference trajectory builder: adapts raw mocap foot positions to the target terrain.

Core outputs: left_precomputed / right_precomputed: (N, 3) left/right ankle reference positions. left_precomputed_quats / right_precomputed_quats: (N, 4) left/right ankle reference quaternions.

Processing pipeline: 1. Create :class:FootstepResolver → resolve a safe landing position for every touchdown frame. 2. :meth:_precompute_foot_reference → stance frames use the locked landing position, swing frames use Hermite/quintic interpolation + floor clamping. 3. :meth:_flatten_stance_quats → the sole stays level during stance (only yaw is kept).

PARAMETER DESCRIPTION
terrain

Terrain object.

TYPE: TerrainSource

motion

Mocap motion data.

TYPE: MotionClip

lookahead

Terrain lookahead distance (m), lets the pelvis lift ahead of upcoming steps.

TYPE: float DEFAULT: 0.0

smoothing_alpha

EMA smoothing coefficient for the pelvis z.

TYPE: float DEFAULT: 0.08

toe_offset_x

Toe x offset relative to the ankle (local frame). Defaults to motion.robot_spec.toe_probe_offset_x.

TYPE: float or None DEFAULT: None

heel_offset_x

Heel x offset relative to the ankle (local frame, negative = behind).

TYPE: float DEFAULT: -0.04

mid_offset_x

Midfoot x offset relative to the ankle (local frame).

TYPE: float DEFAULT: 0.1

swing_floor_margin

Extra ground-clearance margin during swing (m).

TYPE: float DEFAULT: 0.003

soft_reach_limit

Soft pelvis-to-foot reach limit (m).

TYPE: float DEFAULT: 0.7

max_landing_shift

Maximum landing-position search displacement along fwd (m).

TYPE: float DEFAULT: 0.05

Source code in eden/extensions/motion_optimization/tcrs/motion.py
def __init__(
    self,
    terrain: TerrainSource,
    motion: MotionClip,
    lookahead: float = 0.0,
    smoothing_alpha: float = 0.08,
    toe_offset_x: float | None = None,
    heel_offset_x: float = -0.04,
    mid_offset_x: float = 0.10,
    swing_floor_margin: float = 0.003,
    soft_reach_limit: float = 0.70,
    max_landing_shift: float = 0.05,
):
    robot_spec = motion.robot_spec
    if toe_offset_x is None:
        toe_offset_x = robot_spec.toe_probe_offset_x
    lfoot_idx = robot_spec.left_foot_index
    rfoot_idx = robot_spec.right_foot_index
    self.terrain = terrain
    self.motion = motion
    self.lookahead = float(lookahead)
    self.alpha = float(smoothing_alpha)
    self.swing_floor_margin = max(float(swing_floor_margin), 0.0)
    self.soft_reach_limit = max(float(soft_reach_limit), 0.1)
    self._ema_z = None  # exponential moving average of the pelvis z

    # Local x offsets of the foot key points: [heel, midfoot, toe].
    # FootstepResolver uses these three points to check whether the whole foot fits on a step.
    foot_point_offsets = np.array([heel_offset_x, mid_offset_x, toe_offset_x], dtype=np.float64)
    self.left_resolver = FootstepResolver(
        robot_spec=robot_spec,
        terrain=terrain,
        raw_foot=motion.links_pos[:, lfoot_idx],
        contact_mask=motion.left_contact,
        foot_nominal_z=motion.foot_nominal_z,
        foot_quats=motion.links_quat[:, lfoot_idx],
        foot_point_offsets=foot_point_offsets,
        toe_offset_x=toe_offset_x,
        max_landing_shift=max_landing_shift,
    )
    self.right_resolver = FootstepResolver(
        robot_spec=robot_spec,
        terrain=terrain,
        raw_foot=motion.links_pos[:, rfoot_idx],
        contact_mask=motion.right_contact,
        foot_nominal_z=motion.foot_nominal_z,
        foot_quats=motion.links_quat[:, rfoot_idx],
        foot_point_offsets=foot_point_offsets,
        toe_offset_x=toe_offset_x,
        max_landing_shift=max_landing_shift,
    )
    # Precompute the full foot reference trajectories
    # (stance frames = locked footholds, swing frames = interpolation).
    self.left_precomputed = self._precompute_foot_reference(self.left_resolver)
    self.right_precomputed = self._precompute_foot_reference(self.right_resolver)
    # Precompute the ankle reference quaternions (initialized to the raw mocap values).
    self.left_precomputed_quats = motion.links_quat[:, lfoot_idx].copy()
    self.right_precomputed_quats = motion.links_quat[:, rfoot_idx].copy()
    # Keep the sole level during stance: extract the yaw angle, drop pitch/roll.
    # On steps the foot should lie flat, not follow any sole tilt present in the mocap.
    self._flatten_stance_quats(self.left_precomputed_quats, motion.left_contact)
    self._flatten_stance_quats(self.right_precomputed_quats, motion.right_contact)

    # Enforce no lateral crossover between left and right foot references.
    # Check in pelvis local frame: left foot must stay on left side.
    self._enforce_no_crossover(motion, min_sep=0.08)
    # Enforce step-order consistency: the cross product of consecutive
    # left/right landing vectors must match the original motion's sign.
    self._enforce_step_order(motion)

Attributes

terrain instance-attribute
terrain = terrain
motion instance-attribute
motion = motion
lookahead instance-attribute
lookahead = float(lookahead)
alpha instance-attribute
alpha = float(smoothing_alpha)
swing_floor_margin instance-attribute
swing_floor_margin = max(float(swing_floor_margin), 0.0)
soft_reach_limit instance-attribute
soft_reach_limit = max(float(soft_reach_limit), 0.1)
left_resolver instance-attribute
left_resolver = FootstepResolver(robot_spec=robot_spec, terrain=terrain, raw_foot=motion.links_pos[:, lfoot_idx], contact_mask=motion.left_contact, foot_nominal_z=motion.foot_nominal_z, foot_quats=motion.links_quat[:, lfoot_idx], foot_point_offsets=foot_point_offsets, toe_offset_x=toe_offset_x, max_landing_shift=max_landing_shift)
right_resolver instance-attribute
right_resolver = FootstepResolver(robot_spec=robot_spec, terrain=terrain, raw_foot=motion.links_pos[:, rfoot_idx], contact_mask=motion.right_contact, foot_nominal_z=motion.foot_nominal_z, foot_quats=motion.links_quat[:, rfoot_idx], foot_point_offsets=foot_point_offsets, toe_offset_x=toe_offset_x, max_landing_shift=max_landing_shift)
left_precomputed instance-attribute
left_precomputed = self._precompute_foot_reference(self.left_resolver)
right_precomputed instance-attribute
right_precomputed = self._precompute_foot_reference(self.right_resolver)
left_precomputed_quats instance-attribute
left_precomputed_quats = motion.links_quat[:, lfoot_idx].copy()
right_precomputed_quats instance-attribute
right_precomputed_quats = motion.links_quat[:, rfoot_idx].copy()

Methods:

reset
reset() -> None

Reset the pelvis-z EMA state.

Source code in eden/extensions/motion_optimization/tcrs/motion.py
def reset(self) -> None:
    """Reset the pelvis-z EMA state."""
    self._ema_z = None
update
update(pelvis_x: float, pelvis_y: float) -> float

Single-frame update of the pelvis-z exponential moving average (EMA).

Formula: ema_z += alpha * (target - ema_z). The smaller alpha, the smoother (but slower) the pelvis z response. For real-time per-frame use (not precomputation).

RETURNS DESCRIPTION
float

The smoothed pelvis z reference for the current frame.

Source code in eden/extensions/motion_optimization/tcrs/motion.py
def update(self, pelvis_x: float, pelvis_y: float) -> float:
    """Single-frame update of the pelvis-z exponential moving average (EMA).

    Formula: ema_z += alpha * (target - ema_z).
    The smaller alpha, the smoother (but slower) the pelvis z response.
    For real-time per-frame use (not precomputation).

    Returns
    -------
    float
        The smoothed pelvis z reference for the current frame.
    """
    target = self._target_height(pelvis_x, pelvis_y)
    if self._ema_z is None:
        self._ema_z = target  # initialize directly on the first frame
    else:
        self._ema_z += self.alpha * (target - self._ema_z)
    return self._ema_z
segment_offsets
segment_offsets(pelvis_positions: ndarray) -> ndarray

Batch-compute EMA z offsets for a sequence of pelvis positions.

PARAMETER DESCRIPTION
pelvis_positions

(N, 3) pelvis position sequence.

TYPE: ndarray

RETURNS DESCRIPTION
ndarray

(N,) EMA-smoothed z reference per frame.

Source code in eden/extensions/motion_optimization/tcrs/motion.py
def segment_offsets(self, pelvis_positions: np.ndarray) -> np.ndarray:
    """Batch-compute EMA z offsets for a sequence of pelvis positions.

    Parameters
    ----------
    pelvis_positions : np.ndarray
        (N, 3) pelvis position sequence.

    Returns
    -------
    np.ndarray
        (N,) EMA-smoothed z reference per frame.
    """
    offsets = np.zeros(pelvis_positions.shape[0], dtype=np.float64)
    ema = 0.0 if self._ema_z is None else float(self._ema_z)
    for i, pelvis in enumerate(pelvis_positions):
        target = self._target_height(pelvis[0], pelvis[1])
        ema += self.alpha * (target - ema)
        offsets[i] = ema
    return offsets
build_segment
build_segment(pelvis: ndarray, lfoot: ndarray, rfoot: ndarray, left_contact: ndarray, right_contact: ndarray, indices: ndarray) -> tuple[ndarray, ndarray, ndarray, ndarray, ndarray]
Source code in eden/extensions/motion_optimization/tcrs/motion.py
def build_segment(
    self,
    pelvis: np.ndarray,
    lfoot: np.ndarray,
    rfoot: np.ndarray,
    left_contact: np.ndarray,
    right_contact: np.ndarray,
    indices: np.ndarray,
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    pelvis_z_offsets = self.segment_offsets(pelvis)

    ref_pelvis = pelvis.copy()
    ref_pelvis[:, 2] += pelvis_z_offsets
    idx = np.asarray(indices, dtype=np.int64)
    ref_lfoot = self.left_precomputed[idx].copy()
    ref_rfoot = self.right_precomputed[idx].copy()
    ref_lfoot_quats = self.left_precomputed_quats[idx].copy()
    ref_rfoot_quats = self.right_precomputed_quats[idx].copy()

    # Layer 1: enforce reachability by lowering pelvis z when either foot
    # exceeds the soft reach limit.  Compute pelvis-to-foot distance and
    # lower pelvis z per horizon step so that both legs stay within L_soft.
    L = self.soft_reach_limit
    for i in range(ref_pelvis.shape[0]):
        for foot in (ref_lfoot[i], ref_rfoot[i]):
            dxy = np.sqrt((ref_pelvis[i, 0] - foot[0]) ** 2 + (ref_pelvis[i, 1] - foot[1]) ** 2)
            # max allowed vertical distance given horizontal offset
            max_dz = np.sqrt(max(L**2 - dxy**2, 0.0))
            # pelvis must be above foot, so cap: pelvis_z <= foot_z + max_dz
            ref_pelvis[i, 2] = min(ref_pelvis[i, 2], foot[2] + max_dz)

    return ref_pelvis, ref_lfoot, ref_rfoot, ref_lfoot_quats, ref_rfoot_quats

Functions:

stabilize_contact_mask

stabilize_contact_mask(mask: ndarray, min_contact_frames: int = 4, min_swing_frames: int = 2) -> ndarray

Remove contact runs shorter than min_contact_frames and fill swing gaps shorter than min_swing_frames that are surrounded by contact on both sides.

Source code in eden/extensions/motion_optimization/tcrs/motion.py
def stabilize_contact_mask(mask: np.ndarray, min_contact_frames: int = 4, min_swing_frames: int = 2) -> np.ndarray:
    """Remove contact runs shorter than ``min_contact_frames`` and fill swing gaps shorter
    than ``min_swing_frames`` that are surrounded by contact on both sides.
    """
    out = np.asarray(mask, dtype=bool).copy()
    n = out.shape[0]
    if n == 0:
        return out

    i = 0
    while i < n:
        j = i + 1
        while j < n and out[j] == out[i]:
            j += 1
        if out[i] and (j - i) < min_contact_frames:
            out[i:j] = False
        i = j

    i = 0
    while i < n:
        j = i + 1
        while j < n and out[j] == out[i]:
            j += 1
        if (not out[i]) and (j - i) < min_swing_frames:
            left_true = (i > 0) and out[i - 1]
            right_true = (j < n) and out[j]
            if left_true and right_true:
                out[i:j] = True
        i = j

    return out

Gait phase

Phase-clock gait analysis for bipedal motion capture data.

Ported from Mondo-Robotics/PMT TCRS/stair_mppi/gait_phase.py @ 88838ca.

Fits a global phase clock φ ∈ [0,1) to periodic locomotion data, with per-foot phase offsets θ_L, θ_R and swing ratios.

Classifies: walking, running, hopping, galloping, standing.

Key design choices: - Period estimation: hybrid of speed autocorrelation + speed minima touchdown detection. Autocorrelation works well for walking; speed minima work for running (where autocorrelation decays too fast). - Phase timeline: local phase resets at detected touchdowns prevent drift from stride frequency variation in long sequences. - Contact z threshold: 0.06 m by default, terrain-invariant. For data where foot z is not ground-referenced, the raw contact mask is empty and the clock is fitted from foot speed alone. - Clips with no significant flight phase are not fitted (fit_phase_clock returns None): z-based contact is reliable for them.

Usage: from eden.extensions.motion_optimization.tcrs.gait_phase import fit_phase_clock, compute_per_frame_phase params = fit_phase_clock(lfoot_z, rfoot_z, lfoot_speed, rfoot_speed, fps) if params is not None: phase = compute_per_frame_phase(params, n_frames)

Classes

GaitPhaseParams dataclass

GaitPhaseParams(period_frames: float, period_sec: float, swing_ratio_left: float, swing_ratio_right: float, phase_offset_left: float, phase_offset_right: float, gait_type: str, fit_quality: float, fps: float, _left_touchdowns: ndarray = None, _right_touchdowns: ndarray = None)

Fitted phase clock parameters for a motion clip.

Attributes

period_frames instance-attribute
period_frames: float
period_sec instance-attribute
period_sec: float
swing_ratio_left instance-attribute
swing_ratio_left: float
swing_ratio_right instance-attribute
swing_ratio_right: float
phase_offset_left instance-attribute
phase_offset_left: float
phase_offset_right instance-attribute
phase_offset_right: float
gait_type instance-attribute
gait_type: str
fit_quality instance-attribute
fit_quality: float
fps instance-attribute
fps: float
has_flight_phase property
has_flight_phase: bool

True if the stance fractions sum to less than one cycle.

Both feet are then airborne at some point of every cycle, whatever the phase offset.

Methods:

PerFramePhase dataclass

PerFramePhase(global_phase: ndarray, left_local_phase: ndarray, right_local_phase: ndarray, left_contact: ndarray, right_contact: ndarray, flight_mask: ndarray, support_weight_left: ndarray, support_weight_right: ndarray)

Per-frame phase information for the full motion clip.

Attributes

global_phase instance-attribute
global_phase: ndarray
left_local_phase instance-attribute
left_local_phase: ndarray
right_local_phase instance-attribute
right_local_phase: ndarray
left_contact instance-attribute
left_contact: ndarray
right_contact instance-attribute
right_contact: ndarray
flight_mask instance-attribute
flight_mask: ndarray
support_weight_left instance-attribute
support_weight_left: ndarray
support_weight_right instance-attribute
support_weight_right: ndarray

Methods:

Functions:

fit_phase_clock

fit_phase_clock(lfoot_z: ndarray, rfoot_z: ndarray, lfoot_speed: ndarray, rfoot_speed: ndarray, fps: float, contact_z_threshold: float = 0.06, contact_speed_threshold: float = 0.5) -> GaitPhaseParams | None

Fit phase clock parameters from motion capture data.

Hybrid period estimation: 1. Try speed minima touchdown detection (robust for running) 2. Fallback to speed autocorrelation (better for slow walking)

PARAMETER DESCRIPTION
lfoot_z

(N,) left foot z positions, m

TYPE: ndarray

rfoot_z

(N,) right foot z positions, m

TYPE: ndarray

lfoot_speed

(N,) left foot speed magnitude, m/s

TYPE: ndarray

rfoot_speed

(N,) right foot speed magnitude, m/s

TYPE: ndarray

fps

frames per second

TYPE: float

contact_z_threshold

foot z (m) below which a foot can be in raw contact

TYPE: float DEFAULT: 0.06

contact_speed_threshold

foot speed (m/s) below which a foot can be in raw contact

TYPE: float DEFAULT: 0.5

RETURNS DESCRIPTION
GaitPhaseParams or None

None if the clip is shorter than 30 frames, has no significant flight phase (both feet off the ground in under 5% of frames by raw contact, where z-based contact is reliable), or is not periodic.

Source code in eden/extensions/motion_optimization/tcrs/gait_phase.py
def fit_phase_clock(
    lfoot_z: np.ndarray,
    rfoot_z: np.ndarray,
    lfoot_speed: np.ndarray,
    rfoot_speed: np.ndarray,
    fps: float,
    contact_z_threshold: float = 0.06,
    contact_speed_threshold: float = 0.5,
) -> GaitPhaseParams | None:
    """Fit phase clock parameters from motion capture data.

    Hybrid period estimation:
    1. Try speed minima touchdown detection (robust for running)
    2. Fallback to speed autocorrelation (better for slow walking)

    Parameters
    ----------
    lfoot_z : np.ndarray
        (N,) left foot z positions, m
    rfoot_z : np.ndarray
        (N,) right foot z positions, m
    lfoot_speed : np.ndarray
        (N,) left foot speed magnitude, m/s
    rfoot_speed : np.ndarray
        (N,) right foot speed magnitude, m/s
    fps : float
        frames per second
    contact_z_threshold : float
        foot z (m) below which a foot can be in raw contact
    contact_speed_threshold : float
        foot speed (m/s) below which a foot can be in raw contact

    Returns
    -------
    GaitPhaseParams or None
        None if the clip is shorter than 30 frames, has no significant flight phase (both
        feet off the ground in under 5% of frames by raw contact, where z-based contact is
        reliable), or is not periodic.
    """
    N = len(lfoot_z)
    if N < 30:
        return None

    # Raw contact masks (z-threshold, no stabilization — preserves short contacts)
    raw_left_contact = (lfoot_z < contact_z_threshold) & (lfoot_speed < contact_speed_threshold)
    raw_right_contact = (rfoot_z < contact_z_threshold) & (rfoot_speed < contact_speed_threshold)

    # Quick check: if raw contact never has both feet simultaneously off ground,
    # the gait is walking-like and z-based detection works fine.
    raw_flight = (~raw_left_contact) & (~raw_right_contact)
    flight_ratio = float(np.mean(raw_flight))
    if flight_ratio < 0.05:
        # Walking gait — no significant flight phase, z-based contact is reliable
        return None

    # Step 1: Detect touchdowns from speed minima
    left_tds = _detect_touchdowns(lfoot_speed)
    right_tds = _detect_touchdowns(rfoot_speed)

    # Step 2: Estimate period — primary: touchdown intervals
    period = _period_from_touchdowns(left_tds)
    period_R = _period_from_touchdowns(right_tds)
    if period is not None and period_R is not None:
        period = (period + period_R) / 2.0
    elif period is None:
        period = period_R

    # Fallback: autocorrelation
    if period is None:
        period = _period_from_autocorrelation(lfoot_speed)
    if period is None:
        period = _period_from_autocorrelation(rfoot_speed)
    if period is None:
        return None

    # Step 3: Phase offset
    # Primary: from touchdown timing; fallback: cross-correlation
    if len(left_tds) >= 2 and len(right_tds) >= 2:
        offset = _estimate_phase_offset_from_touchdowns(left_tds, right_tds, period)
    else:
        offset = _estimate_phase_offset_from_xcorr(lfoot_speed, rfoot_speed, period)

    # Step 4: Swing ratios (from z-based contact, reliable on flat ground)
    swing_ratio_L = _estimate_swing_ratio(raw_left_contact, left_tds, period)
    swing_ratio_R = _estimate_swing_ratio(raw_right_contact, right_tds, period)

    # Step 5: Build params and classify
    params = GaitPhaseParams(
        period_frames=period,
        period_sec=period / fps,
        swing_ratio_left=swing_ratio_L,
        swing_ratio_right=swing_ratio_R,
        phase_offset_left=0.0,
        phase_offset_right=offset,
        gait_type="",
        fit_quality=0.0,
        fps=fps,
        _left_touchdowns=left_tds,
        _right_touchdowns=right_tds,
    )

    offset_diff = min(offset, 1.0 - offset)
    avg_swing = (swing_ratio_L + swing_ratio_R) / 2.0

    if avg_swing < 0.05:
        gait_type = "stand"
    elif offset_diff > 0.35:
        gait_type = "run" if params.has_flight_phase else "walk"
    elif offset_diff < 0.1:
        gait_type = "hop"
    else:
        gait_type = "gallop"
    params.gait_type = gait_type

    # Step 6: Compute fit quality
    phase = compute_per_frame_phase(params, N)
    # If raw contact has any frames, compare
    if np.any(raw_left_contact) or np.any(raw_right_contact):
        agree_L = np.mean(phase.left_contact == raw_left_contact) if np.any(raw_left_contact) else 0.5
        agree_R = np.mean(phase.right_contact == raw_right_contact) if np.any(raw_right_contact) else 0.5
        params.fit_quality = (agree_L + agree_R) / 2.0
    else:
        # No raw contact frames (foot z not ground-referenced): nothing to compare
        # against, so a nominal quality.
        params.fit_quality = 0.7

    return params

compute_per_frame_phase

compute_per_frame_phase(params: GaitPhaseParams, n_frames: int) -> PerFramePhase

Compute per-frame phase, contact masks, and support weights.

Uses detected touchdowns for piecewise linear phase (drift-resistant).

Source code in eden/extensions/motion_optimization/tcrs/gait_phase.py
def compute_per_frame_phase(
    params: GaitPhaseParams,
    n_frames: int,
) -> PerFramePhase:
    """Compute per-frame phase, contact masks, and support weights.

    Uses detected touchdowns for piecewise linear phase (drift-resistant).
    """
    # Build global phase timeline from left foot touchdowns
    left_tds = params._left_touchdowns if params._left_touchdowns is not None else np.array([], dtype=int)
    global_phase = _build_phase_timeline(n_frames, params.period_frames, left_tds)

    # Local phases
    left_local = (global_phase + params.phase_offset_left) % 1.0
    right_local = (global_phase + params.phase_offset_right) % 1.0

    # Contact from phase: stance = local_phase < stance_ratio
    stance_ratio_L = 1.0 - params.swing_ratio_left
    stance_ratio_R = 1.0 - params.swing_ratio_right
    left_contact = left_local < stance_ratio_L
    right_contact = right_local < stance_ratio_R

    flight_mask = (~left_contact) & (~right_contact)

    # Continuous support weights
    support_weight_left = np.ones(n_frames, dtype=np.float64)
    support_weight_right = np.ones(n_frames, dtype=np.float64)

    # Left foot: cosine decay during swing
    swing_mask_L = ~left_contact
    if params.swing_ratio_left > 1e-6:
        swing_progress_L = (left_local[swing_mask_L] - stance_ratio_L) / params.swing_ratio_left
        swing_progress_L = np.clip(swing_progress_L, 0.0, 1.0)
        support_weight_left[swing_mask_L] = 0.5 * (1.0 + np.cos(np.pi * swing_progress_L))

    # Right foot: cosine decay during swing
    swing_mask_R = ~right_contact
    if params.swing_ratio_right > 1e-6:
        swing_progress_R = (right_local[swing_mask_R] - stance_ratio_R) / params.swing_ratio_right
        swing_progress_R = np.clip(swing_progress_R, 0.0, 1.0)
        support_weight_right[swing_mask_R] = 0.5 * (1.0 + np.cos(np.pi * swing_progress_R))

    return PerFramePhase(
        global_phase=global_phase,
        left_local_phase=left_local,
        right_local_phase=right_local,
        left_contact=left_contact,
        right_contact=right_contact,
        flight_mask=flight_mask,
        support_weight_left=support_weight_left,
        support_weight_right=support_weight_right,
    )

Foot-swing optimization (MPPI)

MPPI (Model Predictive Path Integral) foot swing trajectory optimizer.

Sampling-based foot swing planner. Uses cubic-spline-parameterized control points with a pre-computed basis matrix for fast batch expansion.

Usage: from eden.extensions.motion_optimization.tcrs.mppi_foot import MppiFootOptimizer, MppiFootParams opt = MppiFootOptimizer(MppiFootParams()) traj_x, traj_y, traj_z = opt.optimize_swing( liftoff_pos, landing_pos, swing_duration, ref_traj, ref_times, terrain )

Also exposes slerp_quaternion, used throughout the planner stack.

Ported from Mondo-Robotics/PMT TCRS/stair_mppi/mppi_foot.py @ 88838ca.

Classes

MppiFootParams dataclass

MppiFootParams(n_knots: int = 0, n_samples: int = 128, n_iterations: int = 10, n_dense: int = 100, temperature: float = 0.1, noise_std_xy: float = 0.02, noise_std_z: float = 0.04, w_track: float = 10.0, w_smooth: float = 10.0, w_terrain: float = 10000.0, w_ceiling: float = 500.0, h_clearance: float = 0.03, h_contact_z: float = 0.0, ceiling_margin: float = 0.3, seed: int = 42, w_stance_ground: float = 0.0, w_swing_clearance_reward: float = 0.0)

Parameters for the MPPI foot planner.

Attributes

n_knots class-attribute instance-attribute
n_knots: int = 0
n_samples class-attribute instance-attribute
n_samples: int = 128
n_iterations class-attribute instance-attribute
n_iterations: int = 10
n_dense class-attribute instance-attribute
n_dense: int = 100
temperature class-attribute instance-attribute
temperature: float = 0.1
noise_std_xy class-attribute instance-attribute
noise_std_xy: float = 0.02
noise_std_z class-attribute instance-attribute
noise_std_z: float = 0.04
w_track class-attribute instance-attribute
w_track: float = 10.0
w_smooth class-attribute instance-attribute
w_smooth: float = 10.0
w_terrain class-attribute instance-attribute
w_terrain: float = 10000.0
w_ceiling class-attribute instance-attribute
w_ceiling: float = 500.0
h_clearance class-attribute instance-attribute
h_clearance: float = 0.03
h_contact_z class-attribute instance-attribute
h_contact_z: float = 0.0
ceiling_margin class-attribute instance-attribute
ceiling_margin: float = 0.3
seed class-attribute instance-attribute
seed: int = 42
w_stance_ground class-attribute instance-attribute
w_stance_ground: float = 0.0
w_swing_clearance_reward class-attribute instance-attribute
w_swing_clearance_reward: float = 0.0

Methods:

EvalWrapper

EvalWrapper(times: ndarray, values: ndarray)

Wraps dense trajectory arrays, providing an .eval(t) interface with linear interpolation at arbitrary times.

Internally stores (n_dense,) time and value arrays; eval() linearly interpolates the value at any time.

Source code in eden/extensions/motion_optimization/tcrs/mppi_foot.py
def __init__(self, times: np.ndarray, values: np.ndarray):
    self.times = np.asarray(times, dtype=np.float64)
    self.values = np.asarray(values, dtype=np.float64)

Attributes

times instance-attribute
times = np.asarray(times, dtype=np.float64)
values instance-attribute
values = np.asarray(values, dtype=np.float64)

Methods:

eval
eval(t: ndarray) -> ndarray

Evaluate at arbitrary time points (linear interpolation).

Source code in eden/extensions/motion_optimization/tcrs/mppi_foot.py
def eval(self, t: np.ndarray) -> np.ndarray:
    """Evaluate at arbitrary time points (linear interpolation)."""
    return np.interp(t, self.times, self.values)

MppiFootOptimizer

MppiFootOptimizer(params: MppiFootParams | None = None)

Optimize a single-foot swing trajectory via MPPI sampling.

Returns three EvalWrapper objects (x/y/z), as consumed by :func:~eden.extensions.motion_optimization.tcrs.pipeline.precompute_mppi_foot.

Source code in eden/extensions/motion_optimization/tcrs/mppi_foot.py
def __init__(self, params: MppiFootParams | None = None):
    self.params = params or MppiFootParams()
    self.rng = np.random.default_rng(self.params.seed)

Attributes

params instance-attribute
params = params or MppiFootParams()
rng instance-attribute
rng = np.random.default_rng(self.params.seed)

Methods:

optimize_swing
optimize_swing(liftoff_pos: ndarray, landing_pos: ndarray, swing_duration: float, ref_traj: ndarray, ref_times: ndarray, terrain: TerrainSource, v_start: ndarray | None = None, v_end: ndarray | None = None, phase_stance_mask: ndarray | None = None, phase_swing_envelope: ndarray | None = None) -> tuple[EvalWrapper, EvalWrapper, EvalWrapper]

Optimize the xyz trajectory of one swing.

v_start/v_end are ignored (the clamped spline boundary conditions fix zero endpoint velocity). The optional phase-aware inputs (None = disabled) are phase_stance_mask, an (n_dense,) bool mask of stance samples along the dense times, and phase_swing_envelope, an (n_dense,) swing intensity in [0, 1].

RETURNS DESCRIPTION
tuple[EvalWrapper, EvalWrapper, EvalWrapper]

(traj_x, traj_y, traj_z), each with an .eval(t) method over [0, swing_duration].

Source code in eden/extensions/motion_optimization/tcrs/mppi_foot.py
def optimize_swing(
    self,
    liftoff_pos: np.ndarray,
    landing_pos: np.ndarray,
    swing_duration: float,
    ref_traj: np.ndarray,
    ref_times: np.ndarray,
    terrain: TerrainSource,
    v_start: np.ndarray | None = None,
    v_end: np.ndarray | None = None,
    phase_stance_mask: np.ndarray | None = None,
    phase_swing_envelope: np.ndarray | None = None,
) -> tuple[EvalWrapper, EvalWrapper, EvalWrapper]:
    """Optimize the xyz trajectory of one swing.

    ``v_start``/``v_end`` are ignored (the clamped spline boundary conditions
    fix zero endpoint velocity). The optional phase-aware inputs (None = disabled)
    are ``phase_stance_mask``, an ``(n_dense,)`` bool mask of stance samples along
    the dense times, and ``phase_swing_envelope``, an ``(n_dense,)`` swing
    intensity in [0, 1].

    Returns
    -------
    tuple[EvalWrapper, EvalWrapper, EvalWrapper]
        ``(traj_x, traj_y, traj_z)``, each with an ``.eval(t)`` method over ``[0, swing_duration]``.
    """
    p = self.params
    n_knots_inner = p.n_knots
    n_knots_total = n_knots_inner + 2  # Including the two boundary points

    # ---- Time setup ----
    dense_times = np.linspace(0.0, swing_duration, p.n_dense)
    knot_times = np.linspace(0.0, swing_duration, n_knots_total)

    # ---- Pre-compute the basis matrix ----
    B_full = _build_basis_matrix(knot_times, dense_times)
    # Separate boundary columns from interior columns
    B_start = B_full[:, 0]  # (n_dense,)
    B_end = B_full[:, -1]  # (n_dense,)
    B_inner = B_full[:, 1:-1]  # (n_dense, n_knots_inner)

    # Boundary contribution: (n_dense, 3) — contribution of the fixed endpoints to the trajectory
    boundary_contrib = B_start[:, None] * liftoff_pos[None, :] + B_end[:, None] * landing_pos[None, :]

    # ---- Interpolate the reference trajectory onto dense_times ----
    ref_interp = np.column_stack(
        [np.interp(dense_times, ref_times, ref_traj[:, d]) for d in range(3)]
    )  # (n_dense, 3)

    # ---- Clearance envelope: sin²(πφ), zero at endpoints, maximal in the middle ----
    phi = dense_times / max(swing_duration, 1e-12)
    swing_envelope = np.sin(np.pi * phi) ** 2
    clearance = p.h_clearance * 2.0 * swing_envelope  # (n_dense,)

    # ---- Height upper bound ----
    z_ceiling = max(liftoff_pos[2], landing_pos[2]) + p.ceiling_margin

    # ---- Initialize mu: linearly interpolate the interior points ----
    mu = np.zeros((n_knots_inner, 3), dtype=np.float64)
    for i in range(n_knots_inner):
        alpha = (i + 1) / (n_knots_inner + 1)
        mu[i] = (1 - alpha) * liftoff_pos + alpha * landing_pos

    # ---- Noise standard deviation per channel ----
    noise_std = np.array([p.noise_std_xy, p.noise_std_xy, p.noise_std_z], dtype=np.float64)

    # ---- MPPI iterations ----
    for _iter in range(p.n_iterations):
        # Sampling: (n_samples, n_knots_inner, 3)
        noise = self.rng.normal(size=(p.n_samples, n_knots_inner, 3))
        noise *= noise_std[None, None, :]
        samples = mu[None, :, :] + noise
        samples[0] = mu  # Sample 0 = current mean (noise-free)

        # Expand to dense trajectories: (n_samples, n_dense, 3)
        # trajs = boundary_contrib + B_inner @ samples_per_channel
        trajs = boundary_contrib[None, :, :] + np.einsum("hk,ikd->ihd", B_inner, samples)

        # Terrain query: (n_samples, n_dense)
        terrain_z = np.empty((p.n_samples, p.n_dense), dtype=np.float64)
        for i in range(p.n_samples):
            terrain_z[i] = terrain.height_batch(trajs[i, :, 0], trajs[i, :, 1])

        # Compute costs
        costs = self._cost_batch(
            trajs,
            terrain_z,
            ref_interp,
            clearance,
            z_ceiling,
            phase_stance_mask=phase_stance_mask,
            phase_swing_envelope=phase_swing_envelope,
        )

        # Softmax weighting
        beta = costs.min()
        weights = np.exp(-(costs - beta) / p.temperature)
        w_sum = weights.sum()
        if w_sum < 1e-30:
            weights = np.ones(p.n_samples) / p.n_samples
        else:
            weights /= w_sum

        # Update mu
        mu = np.einsum("i,ijk->jk", weights, samples)

    # ---- Final trajectory ----
    final_traj = boundary_contrib + B_inner @ mu  # (n_dense, 3)

    return (
        EvalWrapper(dense_times, final_traj[:, 0]),
        EvalWrapper(dense_times, final_traj[:, 1]),
        EvalWrapper(dense_times, final_traj[:, 2]),
    )

Functions:

slerp_quaternion

slerp_quaternion(q0_wxyz: ndarray, q1_wxyz: ndarray, t: ndarray) -> ndarray

SLERP between two [w,x,y,z] quaternions. Returns (N, 4) in [w,x,y,z].

Aligns q1 to the same hemisphere as q0 (flips if dot < 0) to ensure shortest-arc interpolation — avoids unexpected 360 degree foot rotations.

Source code in eden/extensions/motion_optimization/tcrs/mppi_foot.py
def slerp_quaternion(q0_wxyz: np.ndarray, q1_wxyz: np.ndarray, t: np.ndarray) -> np.ndarray:
    """SLERP between two [w,x,y,z] quaternions. Returns (N, 4) in [w,x,y,z].

    Aligns q1 to the same hemisphere as q0 (flips if dot < 0) to ensure
    shortest-arc interpolation — avoids unexpected 360 degree foot rotations.
    """
    t = np.asarray(t, dtype=np.float64)
    q0 = np.asarray(q0_wxyz, dtype=np.float64)
    q1 = np.asarray(q1_wxyz, dtype=np.float64)
    if np.dot(q0, q1) < 0.0:
        q1 = -q1
    q0_xyzw = np.array([q0[1], q0[2], q0[3], q0[0]])
    q1_xyzw = np.array([q1[1], q1[2], q1[3], q1[0]])
    rots = Rotation.from_quat(np.stack([q0_xyzw, q1_xyzw]))
    slerp_obj = Slerp([0.0, 1.0], rots)
    interp = slerp_obj(np.clip(t, 0.0, 1.0))
    q_xyzw = interp.as_quat()
    return np.column_stack([q_xyzw[:, 3], q_xyzw[:, 0], q_xyzw[:, 1], q_xyzw[:, 2]])

Pelvis-height filtering

Support-aware pelvis-z filtering for terrain reference warping.

Smooths a per-frame absolute pelvis-z target with an asymmetric EMA and a per-frame rate limit, so the root tracks terrain height changes without step discontinuities. The caller (the pipeline's pelvis-z retargeting stage) constructs the target; this filter only smooths it.

Ported from Mondo-Robotics/PMT TCRS/stair_mppi/terrain_warp.py @ 88838ca (the SupportAwareRootZFilter.step path only; upstream's legacy TerrainReferenceWarper / raw-foot-data filter path is dead code and was not ported).

Classes

SupportAwareRootZFilter

SupportAwareRootZFilter(alpha_up: float = 0.35, alpha_down: float = 0.1, max_delta: float = 0.02, flight_alpha: float = 0.6)

Smooth a precomputed absolute pelvis-z target with EMA + rate limit.

The caller builds z_target externally; this filter applies an asymmetric EMA (faster rise than fall, faster still in flight) followed by a per-frame rate clamp.

Source code in eden/extensions/motion_optimization/tcrs/terrain_warp.py
def __init__(
    self,
    alpha_up: float = 0.35,
    alpha_down: float = 0.10,
    max_delta: float = 0.02,
    flight_alpha: float = 0.6,
):
    self.alpha_up = alpha_up
    self.alpha_down = alpha_down
    self.max_delta = max_delta
    self.flight_alpha = flight_alpha
    self._ema_z = None

Attributes

alpha_up instance-attribute
alpha_up = alpha_up
alpha_down instance-attribute
alpha_down = alpha_down
max_delta instance-attribute
max_delta = max_delta
flight_alpha instance-attribute
flight_alpha = flight_alpha

Methods:

reset
reset()
Source code in eden/extensions/motion_optimization/tcrs/terrain_warp.py
def reset(self):
    self._ema_z = None
step
step(z_target: float, is_flight: bool = False) -> float

Smooth one externally-constructed z_target.

PARAMETER DESCRIPTION
z_target

Absolute pelvis z target for this frame.

TYPE: float

is_flight

If True, track with at least flight_alpha.

TYPE: bool DEFAULT: False

RETURNS DESCRIPTION
float

Smoothed pelvis z.

Source code in eden/extensions/motion_optimization/tcrs/terrain_warp.py
def step(self, z_target: float, is_flight: bool = False) -> float:
    """Smooth *one* externally-constructed z_target.

    Parameters
    ----------
    z_target : float
        Absolute pelvis z target for this frame.
    is_flight : bool
        If True, track with at least ``flight_alpha``.

    Returns
    -------
    float
        Smoothed pelvis z.
    """
    if self._ema_z is None:
        self._ema_z = float(z_target)
        return self._ema_z

    alpha = self.alpha_up if z_target > self._ema_z else self.alpha_down
    if is_flight:
        alpha = max(alpha, self.flight_alpha)

    ema_target = self._ema_z + alpha * (z_target - self._ema_z)
    dz = min(max(ema_target - self._ema_z, -self.max_delta), self.max_delta)
    self._ema_z = self._ema_z + dz
    return self._ema_z

Whole-body IK

MuJoCo damped-least-squares whole-body IK on the leg joints.

The root pose is fixed (or softly tracked in :meth:JacobianIK.solve_with_root) to the target root pose. Upper-body joints are held to the provided reference joint angles. Only the leg joints are updated. Foot targets are tracked as multi-point constraints (ankle body plus virtual toe and heel points offset in the ankle frame), with asymmetric penetration weighting and a deterministic multi-seed restart for hard frames.

Robot geometry (link/DOF names, toe/heel contact offsets) comes from :class:~eden.options.extensions.motion_optimization.TCRSRobotOptions; nothing is hardcoded to the G1 beyond the defaults that spec carries.

Ported from Mondo-Robotics/PMT TCRS/stair_mppi/ghost_ik.py @ 88838ca (class G1GhostJacobianIK; config loading takes the dict from TCRSRobotOptions.ik_config instead of a JSON path).

Classes

IKResult dataclass

IKResult(qpos: ndarray, leg_qpos: ndarray, left_pos_err: float, right_pos_err: float, left_rot_err: float, right_rot_err: float, converged: bool, iterations: int, root_pos_offset: ndarray = None, cost: float = 0.0, seed_index: int = 0, used_multiseed: bool = False)

Attributes

qpos instance-attribute
qpos: ndarray
leg_qpos instance-attribute
leg_qpos: ndarray
left_pos_err instance-attribute
left_pos_err: float
right_pos_err instance-attribute
right_pos_err: float
left_rot_err instance-attribute
left_rot_err: float
right_rot_err instance-attribute
right_rot_err: float
converged instance-attribute
converged: bool
iterations instance-attribute
iterations: int
root_pos_offset class-attribute instance-attribute
root_pos_offset: ndarray = None
cost class-attribute instance-attribute
cost: float = 0.0
seed_index class-attribute instance-attribute
seed_index: int = 0
used_multiseed class-attribute instance-attribute
used_multiseed: bool = False

Methods:

FrameIKTargets dataclass

FrameIKTargets(root_pos: ndarray, root_quat: ndarray, left_pos: ndarray, left_quat: ndarray, right_pos: ndarray, right_quat: ndarray, fixed_upper_body_qpos: ndarray, ref_leg_qpos: ndarray = None, left_point_scale: ndarray = None, right_point_scale: ndarray = None, root_xy_weight: float = 2.0, root_z_weight: float = 10.0)

Per-frame IK targets for trajectory optimization.

Attributes

root_pos instance-attribute
root_pos: ndarray
root_quat instance-attribute
root_quat: ndarray
left_pos instance-attribute
left_pos: ndarray
left_quat instance-attribute
left_quat: ndarray
right_pos instance-attribute
right_pos: ndarray
right_quat instance-attribute
right_quat: ndarray
fixed_upper_body_qpos instance-attribute
fixed_upper_body_qpos: ndarray
ref_leg_qpos class-attribute instance-attribute
ref_leg_qpos: ndarray = None
left_point_scale class-attribute instance-attribute
left_point_scale: ndarray = None
right_point_scale class-attribute instance-attribute
right_point_scale: ndarray = None
root_xy_weight class-attribute instance-attribute
root_xy_weight: float = 2.0
root_z_weight class-attribute instance-attribute
root_z_weight: float = 10.0

Methods:

JacobianIK

JacobianIK(model: MjModel, robot_spec: TCRSRobotOptions, max_iters: int = 30, tol: float = 0.001, damping: float = 0.05, step_size: float = 0.5, orientation_weight: float = 0.2, posture_weight: float = 0.02, max_joint_step: float = 0.2, penetration_weight: float = 10.0, ankle_weight: float = 1.0, toe_weight: float = 1.0, heel_weight: float = 1.0, ankle_rot_weight: float = 0.0, multiseed_enabled: bool = True, multiseed_num_seeds: int = 3, multiseed_trigger_pos_err: float = 0.2, multiseed_trigger_joint_step: float = 0.8, multiseed_seed_span: float = 0.35, multiseed_continuity_weight: float = 0.25, multiseed_posture_weight: float = 0.02)

Differential whole-body IK on the leg joints.

Supports per-link position and rotation constraints loaded from an ik_config dict (see :data:~eden.options.extensions.motion_optimization.DEFAULT_IK_CONFIG). Each link in the config can have: - pos_weight: weight for 3D position tracking (0 = disabled) - rot_weight: weight for 3D orientation tracking (0 = disabled)

Special entries (left_toe, left_heel, right_toe, right_heel) define virtual constraint points on the foot bodies with offsets.

Source code in eden/extensions/motion_optimization/tcrs/ik.py
def __init__(
    self,
    model: mujoco.MjModel,
    robot_spec: TCRSRobotOptions,
    max_iters: int = 30,
    tol: float = 1e-3,
    damping: float = 0.05,
    step_size: float = 0.5,
    orientation_weight: float = 0.2,
    posture_weight: float = 0.02,
    max_joint_step: float = 0.20,
    penetration_weight: float = 10.0,
    ankle_weight: float = 1.0,
    toe_weight: float = 1.0,
    heel_weight: float = 1.0,
    ankle_rot_weight: float = 0.0,
    multiseed_enabled: bool = True,
    multiseed_num_seeds: int = 3,
    multiseed_trigger_pos_err: float = 0.20,
    multiseed_trigger_joint_step: float = 0.80,
    multiseed_seed_span: float = 0.35,
    multiseed_continuity_weight: float = 0.25,
    multiseed_posture_weight: float = 0.02,
):
    self.model = model
    self.robot_spec = robot_spec
    # Foot contact geometry offsets in the foot (ankle-roll) body frame.
    self._toe_offset = np.asarray(robot_spec.toe_offset, dtype=np.float64)
    self._heel_offset = np.asarray(robot_spec.heel_offset, dtype=np.float64)
    # Bodies that are part of the foot multi-point system (handled separately
    # from the generic extra body constraints).
    self._foot_multipoint_keys = {
        robot_spec.left_foot_link,
        robot_spec.right_foot_link,
        "left_toe",
        "left_heel",
        "right_toe",
        "right_heel",
    }
    self.max_iters = max(1, int(max_iters))
    self.tol = float(tol)
    self.damping = float(damping)
    self.step_size = float(step_size)
    self.orientation_weight = float(orientation_weight)
    self.posture_weight = max(float(posture_weight), 0.0)
    self.max_joint_step = max(float(max_joint_step), 1e-6)
    self.penetration_weight = max(float(penetration_weight), 1.0)
    self.multiseed_enabled = bool(multiseed_enabled)
    self.multiseed_num_seeds = max(1, int(multiseed_num_seeds))
    self.multiseed_trigger_pos_err = max(float(multiseed_trigger_pos_err), 0.0)
    self.multiseed_trigger_joint_step = max(float(multiseed_trigger_joint_step), 0.0)
    self.multiseed_seed_span = max(float(multiseed_seed_span), 0.0)
    self.multiseed_continuity_weight = max(float(multiseed_continuity_weight), 0.0)
    self.multiseed_posture_weight = max(float(multiseed_posture_weight), 0.0)
    # Per-point weights for multi-point IK: [ankle, toe, heel] per foot.
    # Applied as diagonal scaling to error and Jacobian rows.
    self.point_weights = np.array(
        [
            ankle_weight,
            ankle_weight,
            ankle_weight,
            toe_weight,
            toe_weight,
            toe_weight,
            heel_weight,
            heel_weight,
            heel_weight,
        ],
        dtype=np.float64,
    )
    # Ankle orientation weight: explicit rotation tracking via axis-angle error.
    # Complements the multi-point (toe/heel) position constraints which only
    # weakly constrain yaw.  Set > 0 to enable (e.g. 6.0).
    self.ankle_rot_weight = max(float(ankle_rot_weight), 0.0)

    # Extra body constraints: list of (body_id, pos_weight, rot_weight, link_name).
    # Populated by from_config() with the bodies outside _foot_multipoint_keys
    # (and other than the pelvis) that have non-zero weights.
    self.body_constraints = []

    # Knee-forward penalty weight (loaded from config solver section).
    # Penalizes configurations where the knee is behind the pelvis
    # in the body's forward (x) direction. Prevents leg-flip solutions.
    self.knee_forward_weight = 0.0

    self.left_foot_body_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, robot_spec.left_foot_link)
    self.right_foot_body_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, robot_spec.right_foot_link)
    if self.left_foot_body_id == -1 or self.right_foot_body_id == -1:
        raise ValueError(
            f"Failed to locate foot bodies '{robot_spec.left_foot_link}'/'{robot_spec.right_foot_link}' "
            "in the MuJoCo model"
        )

    # Knee body IDs for knee-forward penalty
    self.left_knee_body_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, robot_spec.left_knee_link)
    self.right_knee_body_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, robot_spec.right_knee_link)
    self.pelvis_body_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, robot_spec.pelvis_link)

    left_joint_ids = np.array(
        [mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, name) for name in robot_spec.left_leg_dofs_name],
        dtype=np.int32,
    )
    right_joint_ids = np.array(
        [mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, name) for name in robot_spec.right_leg_dofs_name],
        dtype=np.int32,
    )
    if np.any(left_joint_ids < 0) or np.any(right_joint_ids < 0):
        missing = [
            name
            for name in (*robot_spec.left_leg_dofs_name, *robot_spec.right_leg_dofs_name)
            if mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, name) < 0
        ]
        raise ValueError(f"Failed to locate leg joint(s) {missing} in the MuJoCo model")

    self.leg_joint_ids = np.concatenate([left_joint_ids, right_joint_ids], axis=0)
    self.leg_qpos_adr = np.array(
        [model.jnt_qposadr[jid] for jid in self.leg_joint_ids],
        dtype=np.int32,
    )
    self.leg_dof_adr = np.array(
        [model.jnt_dofadr[jid] for jid in self.leg_joint_ids],
        dtype=np.int32,
    )
    self.leg_limits = model.jnt_range[self.leg_joint_ids].copy()
    # The model's limits are kept as they are (mechanical extremes); non-physical solutions
    # are guarded against dynamically in solve_with_root (max_joint_step clamping, soft-limit
    # damping, the multi-seed continuity cost). The margin only keeps seeds strictly inside.
    self._leg_limit_margin = 1e-3 * (self.leg_limits[:, 1] - self.leg_limits[:, 0])
    self._effective_multiseed_num = max(self.multiseed_num_seeds, 5)
    self._halton_seed_offsets = self._build_halton_offsets(
        max(self._effective_multiseed_num, 1), len(self.leg_qpos_adr)
    )

    all_joint_qpos_adr = np.arange(7, model.nq, dtype=np.int32)
    self.non_leg_qpos_adr = np.setdiff1d(all_joint_qpos_adr, self.leg_qpos_adr)
    self.non_leg_local_idx = self.non_leg_qpos_adr - 7

    # Reusable Jacobian buffers
    self._jacp_left = np.zeros((3, model.nv), dtype=np.float64)
    self._jacr_left = np.zeros((3, model.nv), dtype=np.float64)
    self._jacp_right = np.zeros((3, model.nv), dtype=np.float64)
    self._jacr_right = np.zeros((3, model.nv), dtype=np.float64)
    self._quat_left = np.zeros(4, dtype=np.float64)
    self._quat_right = np.zeros(4, dtype=np.float64)
    self._err_rot_left = np.zeros(3, dtype=np.float64)
    self._err_rot_right = np.zeros(3, dtype=np.float64)

    # Multi-point IK buffers: toe and heel Jacobians for each foot
    self._jacp_left_toe = np.zeros((3, model.nv), dtype=np.float64)
    self._jacp_left_heel = np.zeros((3, model.nv), dtype=np.float64)
    self._jacp_right_toe = np.zeros((3, model.nv), dtype=np.float64)
    self._jacp_right_heel = np.zeros((3, model.nv), dtype=np.float64)
    # MuJoCo point buffer (3,) for mj_jac calls
    self._point_buf = np.zeros(3, dtype=np.float64)

    # Generic body constraint buffers (allocated on first use or by from_config)
    self._jacp_body = np.zeros((3, model.nv), dtype=np.float64)
    self._jacr_body = np.zeros((3, model.nv), dtype=np.float64)
    self._quat_err_buf = np.zeros(4, dtype=np.float64)
    self._rot_err_buf = np.zeros(3, dtype=np.float64)

    # Second MjData for computing reference body poses via FK
    self._ref_data = mujoco.MjData(model)

Attributes

model instance-attribute
model = model
robot_spec instance-attribute
robot_spec = robot_spec
max_iters instance-attribute
max_iters = max(1, int(max_iters))
tol instance-attribute
tol = float(tol)
damping instance-attribute
damping = float(damping)
step_size instance-attribute
step_size = float(step_size)
orientation_weight instance-attribute
orientation_weight = float(orientation_weight)
posture_weight instance-attribute
posture_weight = max(float(posture_weight), 0.0)
max_joint_step instance-attribute
max_joint_step = max(float(max_joint_step), 1e-06)
penetration_weight instance-attribute
penetration_weight = max(float(penetration_weight), 1.0)
multiseed_enabled instance-attribute
multiseed_enabled = bool(multiseed_enabled)
multiseed_num_seeds instance-attribute
multiseed_num_seeds = max(1, int(multiseed_num_seeds))
multiseed_trigger_pos_err instance-attribute
multiseed_trigger_pos_err = max(float(multiseed_trigger_pos_err), 0.0)
multiseed_trigger_joint_step instance-attribute
multiseed_trigger_joint_step = max(float(multiseed_trigger_joint_step), 0.0)
multiseed_seed_span instance-attribute
multiseed_seed_span = max(float(multiseed_seed_span), 0.0)
multiseed_continuity_weight instance-attribute
multiseed_continuity_weight = max(float(multiseed_continuity_weight), 0.0)
multiseed_posture_weight instance-attribute
multiseed_posture_weight = max(float(multiseed_posture_weight), 0.0)
point_weights instance-attribute
point_weights = np.array([ankle_weight, ankle_weight, ankle_weight, toe_weight, toe_weight, toe_weight, heel_weight, heel_weight, heel_weight], dtype=np.float64)
ankle_rot_weight instance-attribute
ankle_rot_weight = max(float(ankle_rot_weight), 0.0)
body_constraints instance-attribute
body_constraints = []
knee_forward_weight instance-attribute
knee_forward_weight = 0.0
left_foot_body_id instance-attribute
left_foot_body_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, robot_spec.left_foot_link)
right_foot_body_id instance-attribute
right_foot_body_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, robot_spec.right_foot_link)
left_knee_body_id instance-attribute
left_knee_body_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, robot_spec.left_knee_link)
right_knee_body_id instance-attribute
right_knee_body_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, robot_spec.right_knee_link)
pelvis_body_id instance-attribute
pelvis_body_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, robot_spec.pelvis_link)
leg_joint_ids instance-attribute
leg_joint_ids = np.concatenate([left_joint_ids, right_joint_ids], axis=0)
leg_qpos_adr instance-attribute
leg_qpos_adr = np.array([model.jnt_qposadr[jid] for jid in self.leg_joint_ids], dtype=np.int32)
leg_dof_adr instance-attribute
leg_dof_adr = np.array([model.jnt_dofadr[jid] for jid in self.leg_joint_ids], dtype=np.int32)
leg_limits instance-attribute
leg_limits = model.jnt_range[self.leg_joint_ids].copy()
non_leg_qpos_adr instance-attribute
non_leg_qpos_adr = np.setdiff1d(all_joint_qpos_adr, self.leg_qpos_adr)
non_leg_local_idx instance-attribute
non_leg_local_idx = self.non_leg_qpos_adr - 7

Methods:

from_config classmethod
from_config(model: MjModel, robot_spec: TCRSRobotOptions) -> 'JacobianIK'

Create the IK solver from robot_spec.ik_config (GMR-style per-link weights).

The config dict has two sections: - link_weights: per-link pos_weight/rot_weight - solver: global solver parameters

Real MuJoCo bodies with non-zero pos/rot weights (excluding the foot multi-point entries) are added as extra body constraints.

The solver section is merged over DEFAULT_IK_CONFIG["solver"] and read without per-key fallbacks — the defaults dict is the single source of truth, and unknown solver keys raise.

Source code in eden/extensions/motion_optimization/tcrs/ik.py
@classmethod
def from_config(cls, model: mujoco.MjModel, robot_spec: TCRSRobotOptions) -> "JacobianIK":
    """Create the IK solver from ``robot_spec.ik_config`` (GMR-style per-link weights).

    The config dict has two sections:
      - link_weights: per-link pos_weight/rot_weight
      - solver: global solver parameters

    Real MuJoCo bodies with non-zero pos/rot weights (excluding the foot
    multi-point entries) are added as extra body constraints.

    The solver section is merged over ``DEFAULT_IK_CONFIG["solver"]`` and
    read without per-key fallbacks — the defaults dict is the single source
    of truth, and unknown solver keys raise.
    """
    cfg = robot_spec.ik_config
    solver_cfg = merged_ik_solver_config(cfg.get("solver", {}))
    link_cfg = cfg.get("link_weights", {})

    # Extract foot multi-point weights from link_weights
    ankle_w = link_cfg.get(robot_spec.left_foot_link, {}).get("pos_weight", 1.0)
    ankle_rw = link_cfg.get(robot_spec.left_foot_link, {}).get("rot_weight", 0.0)
    toe_w = link_cfg.get("left_toe", {}).get("pos_weight", 1.0)
    heel_w = link_cfg.get("left_heel", {}).get("pos_weight", 1.0)

    instance = cls(
        model=model,
        robot_spec=robot_spec,
        max_iters=solver_cfg["max_iters"],
        tol=solver_cfg["tol"],
        damping=solver_cfg["damping"],
        step_size=solver_cfg["step_size"],
        orientation_weight=solver_cfg["orientation_weight"],
        posture_weight=solver_cfg["posture_weight"],
        max_joint_step=solver_cfg["max_joint_step"],
        penetration_weight=solver_cfg["penetration_weight"],
        ankle_weight=ankle_w,
        toe_weight=toe_w,
        heel_weight=heel_w,
        ankle_rot_weight=ankle_rw,
        multiseed_enabled=solver_cfg["multiseed_enabled"],
        multiseed_num_seeds=solver_cfg["multiseed_num_seeds"],
        multiseed_trigger_pos_err=solver_cfg["multiseed_trigger_pos_err"],
        multiseed_trigger_joint_step=solver_cfg["multiseed_trigger_joint_step"],
        multiseed_seed_span=solver_cfg["multiseed_seed_span"],
        multiseed_continuity_weight=solver_cfg["multiseed_continuity_weight"],
        multiseed_posture_weight=solver_cfg["multiseed_posture_weight"],
    )

    # Collect extra body constraints (non-foot-multipoint links)
    for link_name, bw in link_cfg.items():
        if link_name in instance._foot_multipoint_keys:
            continue
        if link_name.startswith("_"):
            continue
        pw = bw.get("pos_weight", 0.0)
        rw = bw.get("rot_weight", 0.0)
        if pw == 0.0 and rw == 0.0:
            continue
        # The pelvis is handled separately via root_xy/z_weight
        if link_name == robot_spec.pelvis_link:
            continue
        bid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, link_name)
        if bid < 0:
            raise ValueError(
                f"ik_config link_weights references '{link_name}', which is not a body in the MuJoCo model."
            )
        instance.body_constraints.append((bid, float(pw), float(rw), link_name))

    if instance.body_constraints:
        names = [bc[3] for bc in instance.body_constraints]
        en.logger.info(f"[IK config] Extra body constraints: {names}")

    # Knee-forward penalty
    kfw = solver_cfg["knee_forward_weight"]
    instance.knee_forward_weight = float(kfw)
    if kfw > 0:
        en.logger.debug(f"[IK config] Knee-forward penalty weight: {kfw}")

    if instance.ankle_rot_weight > 0:
        en.logger.debug(f"[IK config] Ankle rotation weight: {instance.ankle_rot_weight}")

    return instance
extract_leg_qpos
extract_leg_qpos(qpos: ndarray) -> ndarray
Source code in eden/extensions/motion_optimization/tcrs/ik.py
def extract_leg_qpos(self, qpos: np.ndarray) -> np.ndarray:
    return np.asarray(qpos[self.leg_qpos_adr], dtype=np.float64).copy()
solve_with_root
solve_with_root(data: MjData, target_root_pos: ndarray, target_root_quat: ndarray, target_lf_pos: ndarray, target_lf_quat: ndarray, target_rf_pos: ndarray, target_rf_quat: ndarray, fixed_upper_body_qpos: ndarray, initial_leg_qpos: ndarray = None, ref_leg_qpos: ndarray = None, root_xy_weight: float = 2.0, root_z_weight: float = 10.0, left_point_scale: ndarray = None, right_point_scale: ndarray = None, allow_multiseed: bool = True, force_multiseed: bool = False, initial_root_pos: ndarray = None, root_continuity_weight: float = 0.0) -> IKResult

Solve IK with cuRobo-style adaptive multi-seed fallback.

The fast path is differential IK from the previous frame seed. If that candidate has high pose error or a large joint jump, run a small deterministic seed bank and select the lowest-cost result.

Source code in eden/extensions/motion_optimization/tcrs/ik.py
def solve_with_root(
    self,
    data: mujoco.MjData,
    target_root_pos: np.ndarray,
    target_root_quat: np.ndarray,
    target_lf_pos: np.ndarray,
    target_lf_quat: np.ndarray,
    target_rf_pos: np.ndarray,
    target_rf_quat: np.ndarray,
    fixed_upper_body_qpos: np.ndarray,
    initial_leg_qpos: np.ndarray = None,
    ref_leg_qpos: np.ndarray = None,
    root_xy_weight: float = 2.0,
    root_z_weight: float = 10.0,
    left_point_scale: np.ndarray = None,
    right_point_scale: np.ndarray = None,
    allow_multiseed: bool = True,
    force_multiseed: bool = False,
    initial_root_pos: np.ndarray = None,
    root_continuity_weight: float = 0.0,
) -> IKResult:
    """Solve IK with cuRobo-style adaptive multi-seed fallback.

    The fast path is differential IK from the previous frame seed.  If
    that candidate has high pose error or a large joint jump, run a small
    deterministic seed bank and select the lowest-cost result.
    """
    primary = self._solve_with_root_single(
        data,
        target_root_pos=target_root_pos,
        target_root_quat=target_root_quat,
        target_lf_pos=target_lf_pos,
        target_lf_quat=target_lf_quat,
        target_rf_pos=target_rf_pos,
        target_rf_quat=target_rf_quat,
        fixed_upper_body_qpos=fixed_upper_body_qpos,
        initial_leg_qpos=initial_leg_qpos,
        ref_leg_qpos=ref_leg_qpos,
        root_xy_weight=root_xy_weight,
        root_z_weight=root_z_weight,
        left_point_scale=left_point_scale,
        right_point_scale=right_point_scale,
    )
    primary.cost = self._ik_result_cost(
        primary,
        initial_leg_qpos,
        ref_leg_qpos,
        initial_root_pos=initial_root_pos,
        root_continuity_weight=root_continuity_weight,
    )

    need_multiseed = force_multiseed or self._needs_multiseed(primary, initial_leg_qpos)
    if (not allow_multiseed) or (not need_multiseed):
        data.qpos[:] = primary.qpos
        mujoco.mj_forward(self.model, data)
        return primary

    best = primary
    best_cost = primary.cost
    primary_seed = None if initial_leg_qpos is None else np.asarray(initial_leg_qpos, dtype=np.float64)
    if initial_leg_qpos is not None:
        candidate = self._solve_with_root_single(
            data,
            target_root_pos=target_root_pos,
            target_root_quat=target_root_quat,
            target_lf_pos=target_lf_pos,
            target_lf_quat=target_lf_quat,
            target_rf_pos=target_rf_pos,
            target_rf_quat=target_rf_quat,
            fixed_upper_body_qpos=fixed_upper_body_qpos,
            initial_leg_qpos=None,
            ref_leg_qpos=ref_leg_qpos,
            root_xy_weight=root_xy_weight,
            root_z_weight=root_z_weight,
            left_point_scale=left_point_scale,
            right_point_scale=right_point_scale,
        )
        candidate.cost = self._ik_result_cost(
            candidate,
            initial_leg_qpos,
            ref_leg_qpos,
            initial_root_pos=initial_root_pos,
            root_continuity_weight=root_continuity_weight,
        )
        candidate.seed_index = -1
        candidate.used_multiseed = True
        if candidate.cost < best_cost:
            best = candidate
            best_cost = candidate.cost
    for seed_idx, seed in enumerate(self._candidate_leg_seeds(initial_leg_qpos, ref_leg_qpos)):
        if primary_seed is not None and float(np.max(np.abs(seed - primary_seed))) < 1e-5:
            continue
        candidate = self._solve_with_root_single(
            data,
            target_root_pos=target_root_pos,
            target_root_quat=target_root_quat,
            target_lf_pos=target_lf_pos,
            target_lf_quat=target_lf_quat,
            target_rf_pos=target_rf_pos,
            target_rf_quat=target_rf_quat,
            fixed_upper_body_qpos=fixed_upper_body_qpos,
            initial_leg_qpos=seed,
            ref_leg_qpos=ref_leg_qpos,
            root_xy_weight=root_xy_weight,
            root_z_weight=root_z_weight,
            left_point_scale=left_point_scale,
            right_point_scale=right_point_scale,
        )
        candidate.cost = self._ik_result_cost(
            candidate,
            initial_leg_qpos,
            ref_leg_qpos,
            initial_root_pos=initial_root_pos,
            root_continuity_weight=root_continuity_weight,
        )
        candidate.seed_index = seed_idx + 1
        candidate.used_multiseed = True
        if candidate.cost < best_cost:
            best = candidate
            best_cost = candidate.cost

    best.cost = best_cost
    best.used_multiseed = True
    data.qpos[:] = best.qpos
    data.qvel[:] = 0.0
    mujoco.mj_forward(self.model, data)
    return best
pack_state
pack_state(root_pos: ndarray, leg_qpos: ndarray) -> ndarray

Pack root_xyz (3) + leg_q (12 for the G1) into a flat (3 + n_leg_dofs,) vector.

Source code in eden/extensions/motion_optimization/tcrs/ik.py
def pack_state(self, root_pos: np.ndarray, leg_qpos: np.ndarray) -> np.ndarray:
    """Pack root_xyz (3) + leg_q (12 for the G1) into a flat (3 + n_leg_dofs,) vector."""
    return np.concatenate(
        [
            np.asarray(root_pos, dtype=np.float64)[:3],
            np.asarray(leg_qpos, dtype=np.float64)[: len(self.leg_qpos_adr)],
        ]
    )
unpack_state
unpack_state(x: ndarray) -> tuple

Unpack (3 + n_leg_dofs,) → (root_pos (3,), leg_qpos (n_leg_dofs,)).

Source code in eden/extensions/motion_optimization/tcrs/ik.py
def unpack_state(self, x: np.ndarray) -> tuple:
    """Unpack (3 + n_leg_dofs,) → (root_pos (3,), leg_qpos (n_leg_dofs,))."""
    x = np.asarray(x, dtype=np.float64)
    return x[:3].copy(), x[3 : 3 + len(self.leg_qpos_adr)].copy()

Functions:

Terrain sources

Terrain height-query sources for the TCRS pipeline.

Upstream TCRS queries terrain through MuJoCo raycasts against a scene XML (RaycastTerrain). Eden replaces that seam with a height-field-backed implementation that can be built from any of Eden's terrain descriptions:

  • :class:~eden.options.entities.TerrainOptions (procedural generator grid),
  • a raw :class:~eden.options.entities.EntityOptions pointing at a 3D mesh file (rasterized to a height field by downward raycasts at construction),
  • an exported height-field .npz (see :mod:eden.utils.terrains.io).

The protocol the pipeline actually consumes is small: height_at / height_batch, plus an optional foot_side_penetration the planner probes with hasattr for side-wall push-out. Upstream's steps list, local_heightmap and half_width_y are dead code and are deliberately not ported.

Height queries interpolate the grid bilinearly. This matches the physical surface: Genesis triangulates the height field into a mesh, so the simulated terrain is piecewise-linear between grid points.

Classes

TerrainSource

Bases: Protocol

What the TCRS pipeline requires of a terrain.

Implementations may additionally provide foot_side_penetration; the pipeline probes for it with hasattr (same contract as upstream).

Methods:

height_at
height_at(x: float, y: float = 0.0) -> float
Source code in eden/extensions/motion_optimization/tcrs/terrain_source.py
def height_at(self, x: float, y: float = 0.0) -> float: ...
height_batch
height_batch(x: ndarray, y: ndarray | None = None) -> ndarray
Source code in eden/extensions/motion_optimization/tcrs/terrain_source.py
def height_batch(self, x: np.ndarray, y: np.ndarray | None = None) -> np.ndarray: ...

HeightFieldTerrain

HeightFieldTerrain(height_field_m: ndarray, horizontal_scale: float, origin: tuple[float, float, float] = (0.0, 0.0, 0.0), oob_height: float = 0.0)

Bilinear height queries over a regular grid, in world coordinates.

Out-of-bounds queries return oob_height (upstream returns 0.0 outside the scanned scene — keep that default) and log a one-time warning with the observed out-of-bounds fraction so a mis-placed clip is visible.

Source code in eden/extensions/motion_optimization/tcrs/terrain_source.py
def __init__(
    self,
    height_field_m: np.ndarray,
    horizontal_scale: float,
    origin: tuple[float, float, float] = (0.0, 0.0, 0.0),
    oob_height: float = 0.0,
):
    data = HeightFieldData(height_field_m=height_field_m, horizontal_scale=horizontal_scale, origin=origin)
    self._grid = data.height_field_m.astype(np.float64) + data.origin[2]
    self._horizontal_scale = data.horizontal_scale
    self._origin_xy = np.array(data.origin[:2], dtype=np.float64)
    self._oob_height = float(oob_height)
    self._oob_warned = False

Attributes

SIDE_MARCH_STEP class-attribute instance-attribute
SIDE_MARCH_STEP = 0.005
bounds property
bounds: tuple[float, float, float, float]

Terrain extent as (x_min, x_max, y_min, y_max).

Methods:

from_terrain_options classmethod
from_terrain_options(options: 'TerrainOptions', oob_height: float = 0.0) -> 'HeightFieldTerrain'

Build from a procedural :class:TerrainOptions (runs the generators).

Source code in eden/extensions/motion_optimization/tcrs/terrain_source.py
@classmethod
def from_terrain_options(cls, options: "TerrainOptions", oob_height: float = 0.0) -> "HeightFieldTerrain":
    """Build from a procedural :class:`TerrainOptions` (runs the generators)."""
    data = height_field_from_terrain_options(options)
    return cls(
        height_field_m=data.height_field_m,
        horizontal_scale=data.horizontal_scale,
        origin=data.origin,
        oob_height=oob_height,
    )
from_file classmethod
from_file(path: str, oob_height: float = 0.0) -> 'HeightFieldTerrain'

Build from an exported height-field .npz (:mod:eden.utils.terrains.io).

Source code in eden/extensions/motion_optimization/tcrs/terrain_source.py
@classmethod
def from_file(cls, path: str, oob_height: float = 0.0) -> "HeightFieldTerrain":
    """Build from an exported height-field ``.npz`` (:mod:`eden.utils.terrains.io`)."""
    data = load_height_field(path)
    return cls(
        height_field_m=data.height_field_m,
        horizontal_scale=data.horizontal_scale,
        origin=data.origin,
        oob_height=oob_height,
    )
from_entity_options classmethod
from_entity_options(options: 'EntityOptions', resolution: float = 0.02, miss_height: float = 0.0, oob_height: float = 0.0) -> 'HeightFieldTerrain'

Rasterize a mesh-file entity to a height field by downward raycasts.

Applies the entity's scale, up/front alignment and default_root_pos/default_root_quat so the grid matches where the entity would sit in a scene. Cells the rays miss (holes, overhangs seen from below) get miss_height — never -inf, which would poison every mean/min in the planner; a miss fraction is logged.

Source code in eden/extensions/motion_optimization/tcrs/terrain_source.py
@classmethod
def from_entity_options(
    cls,
    options: "EntityOptions",
    resolution: float = 0.02,
    miss_height: float = 0.0,
    oob_height: float = 0.0,
) -> "HeightFieldTerrain":
    """Rasterize a mesh-file entity to a height field by downward raycasts.

    Applies the entity's ``scale``, ``up``/``front`` alignment and
    ``default_root_pos``/``default_root_quat`` so the grid matches where
    the entity would sit in a scene. Cells the rays miss (holes, overhangs
    seen from below) get ``miss_height`` — never ``-inf``, which would
    poison every mean/min in the planner; a miss fraction is logged.
    """
    import trimesh
    from genesis.utils.geom import quat_to_R

    from eden.utils.geom import align_up_and_front

    if resolution <= 0:
        raise ValueError(f"resolution must be > 0; got {resolution}.")
    file = options.fetch_asset()
    mesh = trimesh.load(file, force="mesh")
    if not isinstance(mesh, trimesh.Trimesh) or mesh.faces.shape[0] == 0:
        raise ValueError(f"'{file}' did not load as a triangle mesh; cannot rasterize a height field from it.")

    # World transform: scale, then up/front alignment, then the user pose —
    # the same composition Entity._create_morph_from_options applies.
    transform = np.eye(4)
    transform[:3, :3] = align_up_and_front(options.up, options.front)
    scale = options.scale if options.scale is not None else 1.0
    mesh.apply_scale(scale)
    mesh.apply_transform(transform)
    pose = np.eye(4)
    pose[:3, :3] = quat_to_R(np.asarray(options.default_root_quat, dtype=np.float64))
    pose[:3, 3] = np.asarray(options.default_root_pos, dtype=np.float64)
    mesh.apply_transform(pose)

    (min_x, min_y, _), (max_x, max_y, max_z) = mesh.bounds
    # Integer cell counts (not np.arange with a float step, which can append a
    # spurious cell past the mesh bounds → an artificial missed border strip).
    nx = int(np.floor((max_x - min_x) / resolution + 1e-9)) + 1
    ny = int(np.floor((max_y - min_y) / resolution + 1e-9)) + 1
    xs = min_x + resolution * np.arange(nx)
    ys = min_y + resolution * np.arange(ny)
    xx, yy = np.meshgrid(xs, ys, indexing="ij")
    n_rays = xx.size
    if n_rays > 4_000_000:
        en.logger.warning(
            f"HeightFieldTerrain.from_entity_options: rasterizing '{file}' casts {n_rays:,} rays "
            f"({nx}x{ny} grid at resolution={resolution} m). This is slow without an accelerated "
            "trimesh ray backend (pyembree/rtree); increase `resolution` if construction stalls."
        )
    ray_origins = np.stack([xx.ravel(), yy.ravel(), np.full(xx.size, max_z + 1.0)], axis=1)
    ray_directions = np.tile([0.0, 0.0, -1.0], (ray_origins.shape[0], 1))

    # multiple_hits=False returns the nearest hit per ray (both trimesh
    # intersectors), which for a from-above downward ray IS the highest —
    # walkable — surface, so no per-ray max reduction is needed.
    locations, index_ray, _ = mesh.ray.intersects_location(ray_origins, ray_directions, multiple_hits=False)
    # Start from -inf so surfaces below miss_height are never clobbered by
    # the fill value; misses are identified by the untouched -inf and only
    # then filled with miss_height.
    height = np.full(xx.size, -np.inf, dtype=np.float64)
    if len(index_ray):
        height[index_ray] = locations[:, 2]
    miss_mask = ~np.isfinite(height)
    height[miss_mask] = float(miss_height)
    n_miss = int(miss_mask.sum())
    if n_miss:
        en.logger.warning(
            f"HeightFieldTerrain.from_entity_options: {n_miss}/{xx.size} raster cells "
            f"({n_miss / xx.size:.1%}) missed '{file}'; filled with miss_height={miss_height}."
        )
    return cls(
        height_field_m=height.reshape(xx.shape),
        horizontal_scale=resolution,
        origin=(float(min_x), float(min_y), 0.0),
        oob_height=oob_height,
    )
state_bytes
state_bytes() -> bytes

Stable byte encoding of everything that determines a height query.

Grid values, cell size, xy origin and the out-of-bounds fill — callers (e.g. the TCRS dataset modifier's result cache) hash these bytes instead of reaching into the private fields.

Source code in eden/extensions/motion_optimization/tcrs/terrain_source.py
def state_bytes(self) -> bytes:
    """Stable byte encoding of everything that determines a height query.

    Grid values, cell size, xy origin and the out-of-bounds fill — callers
    (e.g. the TCRS dataset modifier's result cache) hash these bytes instead
    of reaching into the private fields.
    """
    return b"".join(
        (
            np.ascontiguousarray(self._grid, dtype=np.float64).tobytes(),
            np.float64(self._horizontal_scale).tobytes(),
            np.asarray(self._origin_xy, dtype=np.float64).tobytes(),
            np.float64(self._oob_height).tobytes(),
        )
    )
height_at
height_at(x: float, y: float = 0.0) -> float
Source code in eden/extensions/motion_optimization/tcrs/terrain_source.py
def height_at(self, x: float, y: float = 0.0) -> float:
    return float(self.height_batch(np.asarray(x, dtype=np.float64), np.asarray(y, dtype=np.float64)))
height_batch
height_batch(x: ndarray, y: ndarray | None = None) -> ndarray
Source code in eden/extensions/motion_optimization/tcrs/terrain_source.py
def height_batch(self, x: np.ndarray, y: np.ndarray | None = None) -> np.ndarray:
    x = np.asarray(x, dtype=np.float64)
    y = np.zeros_like(x) if y is None else np.asarray(y, dtype=np.float64)

    i = (x - self._origin_xy[0]) / self._horizontal_scale
    j = (y - self._origin_xy[1]) / self._horizontal_scale
    h, w = self._grid.shape
    in_bounds = (i >= 0) & (i <= h - 1) & (j >= 0) & (j <= w - 1)
    if not in_bounds.all() and not self._oob_warned:
        self._oob_warned = True
        frac = 1.0 - float(np.mean(in_bounds))
        en.logger.warning(
            f"HeightFieldTerrain: {frac:.1%} of height queries fell outside the terrain bounds "
            f"{self.bounds}; returning oob_height={self._oob_height} for those. "
            "Check the clip's terrain placement (terrain offset / origin)."
        )

    i = np.clip(i, 0.0, h - 1)
    j = np.clip(j, 0.0, w - 1)
    # Bilinear interpolation on the pre-clipped coords (probed equivalent to
    # the manual corner math to < 1e-12, incl. h==1/w==1 grids). ravel/reshape
    # keeps 0-d (height_at) queries working.
    z = map_coordinates(self._grid, [i.ravel(), j.ravel()], order=1, mode="nearest").reshape(i.shape)
    return np.where(in_bounds, z, self._oob_height)
print_info
print_info() -> None

Log a terrain summary (upstream prints; Eden routes to the logger).

Source code in eden/extensions/motion_optimization/tcrs/terrain_source.py
def print_info(self) -> None:
    """Log a terrain summary (upstream prints; Eden routes to the logger)."""
    x_min, x_max, y_min, y_max = self.bounds
    en.logger.info(
        f"HeightFieldTerrain: grid {self._grid.shape[0]}x{self._grid.shape[1]} @ "
        f"{self._horizontal_scale} m, x=[{x_min:.2f}, {x_max:.2f}], y=[{y_min:.2f}, {y_max:.2f}], "
        f"z=[{self._grid.min():.3f}, {self._grid.max():.3f}] m"
    )
side_clearance
side_clearance(x: float, y: float, z: float, direction: ndarray, max_dist: float = 2.0) -> float

Distance from (x, y, z) along horizontal direction to the first terrain wall.

A "wall" is where the terrain surface rises above z. Marches the height field in :data:SIDE_MARCH_STEP increments (upstream shoots a horizontal mj_ray); returns max_dist when nothing is hit.

Source code in eden/extensions/motion_optimization/tcrs/terrain_source.py
def side_clearance(self, x: float, y: float, z: float, direction: np.ndarray, max_dist: float = 2.0) -> float:
    """Distance from ``(x, y, z)`` along horizontal ``direction`` to the first terrain wall.

    A "wall" is where the terrain surface rises above ``z``. Marches the
    height field in :data:`SIDE_MARCH_STEP` increments (upstream shoots a
    horizontal ``mj_ray``); returns ``max_dist`` when nothing is hit.
    """
    dir_xy = np.array([float(direction[0]), float(direction[1])], dtype=np.float64)
    norm = np.linalg.norm(dir_xy)
    if norm < 1e-8:
        return max_dist
    dir_xy /= norm
    ts = np.arange(self.SIDE_MARCH_STEP, max_dist + self.SIDE_MARCH_STEP, self.SIDE_MARCH_STEP)
    heights = self.height_batch(x + ts * dir_xy[0], y + ts * dir_xy[1])
    hits = np.nonzero(heights > z)[0]
    if hits.size == 0:
        return max_dist
    return float(ts[hits[0]])
foot_side_penetration
foot_side_penetration(ankle_x: float, ankle_y: float, ankle_z: float, fwd_x: float, fwd_y: float, foot_half_width: float = 0.05, foot_half_length: float = 0.12) -> tuple[float, float]

Push-out correction if the foot rectangle penetrates terrain side walls.

Port of upstream RaycastTerrain.foot_side_penetration: probes 4 horizontal directions at sole height and accumulates a push away from any wall closer than the foot half-extent. Returns (0, 0) if clear.

Source code in eden/extensions/motion_optimization/tcrs/terrain_source.py
def foot_side_penetration(
    self,
    ankle_x: float,
    ankle_y: float,
    ankle_z: float,
    fwd_x: float,
    fwd_y: float,
    foot_half_width: float = 0.05,
    foot_half_length: float = 0.12,
) -> tuple[float, float]:
    """Push-out correction if the foot rectangle penetrates terrain side walls.

    Port of upstream ``RaycastTerrain.foot_side_penetration``: probes 4
    horizontal directions at sole height and accumulates a push away from
    any wall closer than the foot half-extent. Returns ``(0, 0)`` if clear.
    """
    # Probe near sole level, but at least 5 mm above the LOCAL surface. Deviation
    # from upstream, whose absolute 0.005 m floor assumes the ground sits at world
    # z=0: identical when the local surface is at z≈0, but on sunken terrain the
    # absolute floor probed at wall height and reported phantom penetrations.
    check_z = max(ankle_z - 0.03, self.height_at(ankle_x, ankle_y) + 0.005)

    lat_x, lat_y = -fwd_y, fwd_x
    push_x, push_y = 0.0, 0.0
    checks = [
        (fwd_x, fwd_y, foot_half_length),
        (-fwd_x, -fwd_y, foot_half_length),
        (lat_x, lat_y, foot_half_width),
        (-lat_x, -lat_y, foot_half_width),
    ]
    for dx, dy, half_size in checks:
        # max_dist=half_size: any clearance >= half_size means "no push", so
        # marching further is wasted work (~17x fewer height samples).
        dist = self.side_clearance(ankle_x, ankle_y, check_z, np.array([dx, dy]), max_dist=half_size)
        if dist < half_size:
            penetration = half_size - dist
            push_x -= penetration * dx
            push_y -= penetration * dy
    return push_x, push_y

Functions:

Synthetic demo clip

Synthetic demo clip for the TCRS examples and tests (numpy-only).

Lives in the package (not the example scripts) so the recipes and tests can share one generator instead of importlib/sys.path hacks on examples/.

Functions:

make_demo_clip

make_demo_clip(num_frames: int = 100, fps: float = 50.0)

Synthetic straight-line G1 walk along +x (Eden dofs/links order).

Feet alternate 1.4 Hz strides with a 6 cm ankle clearance on flat ground; hips/knees are roughly interpolated so leg-geometry helpers see sane capsule lengths. Returns (fps, dofs_pos, links_pos_w, links_quat_w).

Source code in eden/extensions/motion_optimization/tcrs/demo.py
def make_demo_clip(num_frames: int = 100, fps: float = 50.0):
    """Synthetic straight-line G1 walk along +x (Eden dofs/links order).

    Feet alternate 1.4 Hz strides with a 6 cm ankle clearance on flat ground;
    hips/knees are roughly interpolated so leg-geometry helpers see sane
    capsule lengths. Returns ``(fps, dofs_pos, links_pos_w, links_quat_w)``.
    """
    from eden.options.robots.unitree_g1 import G1_DOFS29_NAME, G1_DOFS29_TRACKED_LINKS_NAME

    dofs, links = list(G1_DOFS29_NAME), list(G1_DOFS29_TRACKED_LINKS_NAME)
    t = np.arange(num_frames) / fps
    phase = 2 * np.pi * 1.4 * t
    pelvis = np.stack([0.45 * t, np.zeros(num_frames), 0.79 + 0.01 * np.sin(2 * phase)], axis=1)

    def foot(sign, off):
        p = phase + off
        swing = np.clip(np.sin(p), 0.0, 1.0)  # > 0 during the swing half of the cycle
        z = 0.06 + 0.08 * swing * np.sin(np.clip(p % (2 * np.pi), 0, np.pi))
        z = np.where(np.sin(p) <= 0, 0.06, z)
        return np.stack([0.45 * t + 0.10 * np.sin(p), np.full(num_frames, sign * 0.12), z], axis=1)

    lfoot, rfoot = foot(+1, 0.0), foot(-1, np.pi)
    links_pos = np.repeat(pelvis[:, None, :], len(links), axis=1)
    links_pos[:, links.index("left_ankle_roll_link")] = lfoot
    links_pos[:, links.index("right_ankle_roll_link")] = rfoot
    for side, f in (("left", lfoot), ("right", rfoot)):
        hip = pelvis + np.array([0.0, (1 if side == "left" else -1) * 0.06, -0.17])
        links_pos[:, links.index(f"{side}_hip_pitch_link")] = hip
        links_pos[:, links.index(f"{side}_knee_link")] = 0.5 * (hip + f) + np.array([0.02, 0.0, 0.0])
    links_quat = np.zeros((num_frames, len(links), 4))
    links_quat[..., 0] = 1.0
    dofs_pos = np.zeros((num_frames, len(dofs)))
    for side, off in (("left", 0.0), ("right", np.pi)):
        dofs_pos[:, dofs.index(f"{side}_knee_joint")] = 0.4 + 0.1 * np.sin(phase + off)
        dofs_pos[:, dofs.index(f"{side}_hip_pitch_joint")] = -0.2 - 0.15 * np.sin(phase + off)
        dofs_pos[:, dofs.index(f"{side}_ankle_pitch_joint")] = -0.2
    return fps, dofs_pos, links_pos, links_quat

See also