pi0_moh.py

import dataclasses
import logging

import einops
import flax.nnx as nnx
import flax.nnx.bridge as nnx_bridge
import jax
import jax.numpy as jnp
from typing_extensions import override
from typing import List, Optional

from openpi.models import model as _model
import openpi.models.gemma as _gemma
import openpi.models.siglip as _siglip
from openpi.shared import array_typing as at
import openpi.shared.nnx_utils as nnx_utils
from openpi.models.pi0moh_config import Pi0GatedConfig


def make_attn_mask(input_mask, mask_ar):
    """Adapted from big_vision.

    Tokens can attend to valid inputs tokens which have a cumulative mask_ar
    smaller or equal to theirs. This way `mask_ar` bool[?B, N] can be used to
    setup several types of attention, for example:

      [[1 1 1 1 1 1]]: pure causal attention.

      [[0 0 0 1 1 1]]: prefix-lm attention. The first 3 tokens can attend between
          themselves and the last 3 tokens have a causal attention. The first
          entry could also be a 1 without changing behaviour.

      [[1 0 1 0 1 0 0 1 0 0]]: causal attention between 4 blocks. Tokens of a
          block can attend all previous blocks and all tokens on the same block.

    Args:
      input_mask: bool[B, N] true if its part of the input, false if padding.
      mask_ar: bool[?B, N] mask that's true where previous tokens cannot depend on
        it and false where it shares the same attention mask as the previous token.
    """
    mask_ar = jnp.broadcast_to(mask_ar, input_mask.shape)
    cumsum = jnp.cumsum(mask_ar, axis=1)
    attn_mask = cumsum[:, None, :] <= cumsum[:, :, None]
    valid_mask = input_mask[:, None, :] * input_mask[:, :, None]
    return jnp.logical_and(attn_mask, valid_mask)


# Copied from pi0.py
@at.typecheck
def posemb_sincos(
    pos: at.Real[at.Array, " b"], embedding_dim: int, min_period: float, max_period: float
) -> at.Float[at.Array, "b {embedding_dim}"]:
    """Computes sine-cosine positional embedding vectors for scalar positions."""
    if embedding_dim % 2 != 0:
        raise ValueError(f"embedding_dim ({embedding_dim}) must be divisible by 2")

    fraction = jnp.linspace(0.0, 1.0, embedding_dim // 2)
    period = min_period * (max_period / min_period) ** fraction
    sinusoid_input = jnp.einsum(
        "i,j->ij",
        pos,
        1.0 / period * 2 * jnp.pi,
        precision=jax.lax.Precision.HIGHEST,
    )
    return jnp.concatenate([jnp.sin(sinusoid_input), jnp.cos(sinusoid_input)], axis=-1)



class Pi0Gated(_model.BaseModel):

    def __init__(self, config: Pi0GatedConfig, rngs: nnx.Rngs):
        # Initialize base model with max_horizon.
        super().__init__(config.action_dim, config.action_horizon, config.max_token_len)
        self.config = config
        self.pi05 = config.pi05

        paligemma_config = _gemma.get_config(config.paligemma_variant)
        action_expert_config = _gemma.get_config(config.action_expert_variant)

        # TODO: rewrite gemma in NNX. For now, use bridge.
        llm = nnx_bridge.ToNNX(
            _gemma.Module(
                configs=[paligemma_config, action_expert_config],
                embed_dtype=config.dtype,
                adarms=config.pi05,
            )
        )
        llm.lazy_init(
            rngs=rngs,
            method="init",
            use_adarms=[False, True] if config.pi05 else [False, False],
        )

        img = nnx_bridge.ToNNX(
            _siglip.Module(
                num_classes=paligemma_config.width,
                variant="So400m/14",
                pool_type="none",
                scan=True,
                dtype_mm=config.dtype,
            )
        )
        img.lazy_init(next(iter(config.fake_obs().images.values())), train=False, rngs=rngs)

        # Attribute names must match pi0.py for weight loading.
        self.PaliGemma = nnx.Dict(llm=llm, img=img)

        # Shared action input projection.
        self.action_in_proj = nnx.Linear(config.action_dim, action_expert_config.width, rngs=rngs)

        if config.pi05:
            # Pi0.5-style: adaRMS conditioning on timestep.
            self.time_mlp_in = nnx.Linear(action_expert_config.width, action_expert_config.width, rngs=rngs)
            self.time_mlp_out = nnx.Linear(action_expert_config.width, action_expert_config.width, rngs=rngs)
        else:
            # Pi0-style: state token + action-time MLP (no adaRMS).
            self.state_proj = nnx.Linear(config.action_dim, action_expert_config.width, rngs=rngs)
            self.action_time_mlp_in = nnx.Linear(
                2 * action_expert_config.width,
                action_expert_config.width,
                rngs=rngs,
            )
            self.action_time_mlp_out = nnx.Linear(
                action_expert_config.width,
                action_expert_config.width,
                rngs=rngs,
            )

        self.action_out_proj = nnx.Linear(action_expert_config.width, config.action_dim, rngs=rngs)

        # Extra gating head for Mixture-of-Horizons.
        self.gate_out_proj = nnx.Linear(action_expert_config.width, 1, rngs=rngs)

    @at.typecheck
    def embed_prefix(
            self, obs: _model.Observation
    ) -> tuple[
        at.Float[at.Array, "b s emb"],
        at.Bool[at.Array, "b s"],
        at.Bool[at.Array, " s"],
    ]:
        """Unchanged from pi0.py"""
        input_mask = []
        ar_mask = []
        tokens = []
        # embed images
        for name in obs.images:
            image_tokens, _ = self.PaliGemma.img(obs.images[name], train=False)
            tokens.append(image_tokens)
            input_mask.append(
                einops.repeat(
                    obs.image_masks[name],
                    "b -> b s",
                    s=image_tokens.shape[1],
                )
            )
            # image tokens attend to each other
            ar_mask += [False] * image_tokens.shape[1]

        # add language (aka tokenized inputs)
        if obs.tokenized_prompt is not None:
            tokenized_inputs = self.PaliGemma.llm(obs.tokenized_prompt, method="embed")
            tokens.append(tokenized_inputs)
            input_mask.append(obs.tokenized_prompt_mask)
            # full attention between image and language inputs
            ar_mask += [False] * tokenized_inputs.shape[1]

        tokens = jnp.concatenate(tokens, axis=1)
        input_mask = jnp.concatenate(input_mask, axis=1)
        ar_mask = jnp.array(ar_mask)
        return tokens, input_mask, ar_mask

    @at.typecheck
    def embed_suffix(
            self, state,
            noisy_actions: _model.Actions,
            timestep: at.Float[at.Array, " b"],
            action_pad_mask,
    ) -> tuple[
        at.Float[at.Array, "b s emb"],
        at.Bool[at.Array, "b s"],
        at.Bool[at.Array, " s"],
        at.Float[at.Array, "b emb"] | None,
    ]:
        """
        Pi0 / Pi0.5 compatible suffix embedding.

        Mirrors :class:`Pi0`'s ``embed_suffix`` (including adaRMS conditioning
        when ``pi05=True``) but takes ``state`` and ``action_pad_mask``
        explicitly to support batched horizon processing used by MoH.
        """
        input_mask = []
        ar_mask: list[bool] = []
        tokens = []

        adarms_cond = None

        # Optional Pi0-style state token (no state token for Pi0.5 / adaRMS).
        if not self.pi05:
            state_token = self.state_proj(state)[:, None, :]
            tokens.append(state_token)
            input_mask.append(jnp.ones((state.shape[0], 1), dtype=jnp.bool_))
            # image/language inputs do not attend to state or actions
            ar_mask += [True]

        # Timestep embedding.
        time_emb = posemb_sincos(
            timestep,
            self.action_in_proj.out_features,
            min_period=4e-3,
            max_period=4.0,
        )

        # Project actions.
        action_tokens = self.action_in_proj(noisy_actions)

        if self.pi05:
            # Pi0.5: adaRMS on time embedding, actions unchanged.
            time_emb = self.time_mlp_in(time_emb)
            time_emb = nnx.swish(time_emb)
            time_emb = self.time_mlp_out(time_emb)
            time_emb = nnx.swish(time_emb)
            adarms_cond = time_emb
            action_expert_tokens = action_tokens
        else:
            # Pi0: concatenate time + actions and pass through MLP.
            time_tokens = einops.repeat(time_emb, "b emb -> b s emb", s=noisy_actions.shape[1])
            action_time_tokens = jnp.concatenate([action_tokens, time_tokens], axis=-1)
            action_time_tokens = self.action_time_mlp_in(action_time_tokens)
            action_time_tokens = nnx.swish(action_time_tokens)
            action_time_tokens = self.action_time_mlp_out(action_time_tokens)
            action_expert_tokens = action_time_tokens

        if action_pad_mask is None:
            action_pad_mask = jnp.ones(action_expert_tokens.shape[:2], dtype=jnp.bool_)
        input_mask.append(action_pad_mask)

        tokens.append(action_expert_tokens)

        # image/language/state inputs do not attend to action tokens
        ar_mask += [True] + ([False] * (action_expert_tokens.shape[1] - 1))

        tokens = jnp.concatenate(tokens, axis=1)
        input_mask_arr = jnp.concatenate(input_mask, axis=1)
        ar_mask_arr = jnp.array(ar_mask)

        return tokens, input_mask_arr, ar_mask_arr, adarms_cond

    def cv_squared(self, x: at.Array, eps: float = 1e-10) -> at.Array:
        """Computes the squared coefficient of variation. (From pi0_pytorch_moh.py)"""

        def compute_cv():
            mean = jnp.mean(x, dtype=jnp.float32)
            var = jnp.var(x, dtype=jnp.float32)
            return var / (jnp.square(mean) + eps)

        # Handle num_experts = 1 case
        return jax.lax.cond(
            x.shape[0] == 1,
            lambda: jnp.array(0.0, dtype=jnp.float32),
            compute_cv
        )

    @override
    def compute_loss(
            self, rng: at.KeyArrayLike, observation: _model.Observation, actions: _model.Actions, *, train: bool = False
    ) -> at.Float[at.Array, "*b ah"]:
    # def compute_loss(
    #     self,
    #     rng: at.KeyArrayLike,
    #     observation: _model.Observation,
    #     actions: at.Float[at.Array, "b s action_dim"],
    # ) -> tuple[at.Float[at.Array, ""], dict[str, at.Float[at.Array, ""]]]:
        preprocess_rng, noise_rng, time_rng = jax.random.split(rng, 3)
        observation = _model.preprocess_observation(preprocess_rng, observation, train=train)

        batch_size, max_horizon, action_dim = actions.shape
        num_horizons = len(self.config.horizons)
        horizons_arr = jnp.array(self.config.horizons)

        # Sample noise and time
        noise = jax.random.normal(noise_rng, actions.shape)
        time_scalar = jax.random.beta(time_rng, 1.5, 1, (batch_size,)) * 0.999 + 0.001

        # Expand time and actions for each horizon
        # time shape: (H, B)
        time = einops.repeat(time_scalar, "b -> h b", h=num_horizons)
        # x_t shape: (H, B, max_H, D)
        x_t = time[..., None, None] * noise[None, ...] + (1 - time[..., None, None]) * actions[None, ...]
        # u_t (target) shape: (B, max_H, D)
        u_t = noise - actions

        # STAGE 1: VLM Prefix Pass (Compute KV cache once)
        prefix_tokens, prefix_mask, prefix_ar_mask = self.embed_prefix(observation)
        prefix_attn_mask = make_attn_mask(prefix_mask, prefix_ar_mask)
        prefix_positions = jnp.cumsum(prefix_mask, axis=1) - 1
        (_, prefix_out), prefix_past_key_values = self.PaliGemma.llm(
            [prefix_tokens, None],
            mask=prefix_attn_mask,
            positions=prefix_positions
        )

        # STAGE 2: Action Head Suffix Passes (Parallelized via batching)

        # Repeat prefix masks and KV cache for each horizon
        # New batch size is (B * H)
        batched_prefix_mask = jnp.repeat(prefix_mask, num_horizons, axis=0)
        batched_past_key_values = jax.tree_map(
            lambda x: jnp.repeat(x, num_horizons, axis=1),
            prefix_past_key_values
        )
        batched_state = jnp.repeat(observation.state, num_horizons, axis=0)

        # Reshape x_t and time to align with the new batch dimension
        # (H, B, max_H, D) -> (B*H, max_H, D)
        batched_x_t = jnp.transpose(x_t, (1, 0, 2, 3)).reshape(batch_size * num_horizons, max_horizon, -1)
        # (H, B) -> (B*H,)
        batched_time = jnp.transpose(time, (1, 0)).reshape(-1)

        # Create a padding mask for actions based on valid horizon length
        # (H, max_H)
        action_pad_mask = jnp.arange(max_horizon)[None, :] < horizons_arr[:, None]
        # (B*H, max_H)
        action_pad_mask_expanded = jnp.broadcast_to(
            action_pad_mask[None, :, :],
            (batch_size, num_horizons, max_horizon)
        )
        # (B, H, max_H) -> (B*H, max_H)
        batched_action_pad_mask = action_pad_mask_expanded.reshape(batch_size * num_horizons, max_horizon)

        # Embed the batched suffix inputs
        suffix_tokens, suffix_pad_masks, suffix_ar_mask, adarms_cond = self.embed_suffix(
            batched_state, batched_x_t, batched_time, action_pad_mask=batched_action_pad_mask
        )

        # Combine prefix and suffix masks for cross-attention
        pad_masks = jnp.concatenate([batched_prefix_mask, suffix_pad_masks], axis=1)
        ar_masks = jnp.concatenate([prefix_ar_mask, suffix_ar_mask], axis=0)
        full_att_2d_masks = make_attn_mask(pad_masks, ar_masks)

        prefix_len = prefix_mask.shape[1]
        suffix_len = suffix_tokens.shape[1]

        # Create position IDs and attention mask for the suffix part only
        suffix_position_ids = jnp.arange(prefix_len, prefix_len + suffix_len)[None, :]
        suffix_att_2d_masks = full_att_2d_masks[:, -suffix_len:, :]

        b = suffix_att_2d_masks.shape[0]
        suffix_position_ids = jnp.broadcast_to(suffix_position_ids, (b, suffix_len))

        adarms = [None, adarms_cond] if self.pi05 else [None, None]
        (_, suffix_out), _ = self.PaliGemma.llm(
            [None, suffix_tokens],
            mask=suffix_att_2d_masks,
            positions=suffix_position_ids,
            kv_cache=batched_past_key_values,
            adarms_cond=adarms,
        )
        
        action_start_index = 0 if self.pi05 else 1  # pi0.5 has no state token
        v_t_batched = self.action_out_proj(suffix_out)
        v_t_actions_padded = v_t_batched[:, action_start_index: action_start_index + max_horizon, :]
        # (H, B, max_H, D_action)
        all_v_t_preds = v_t_actions_padded.reshape(
            batch_size, num_horizons, max_horizon, -1
        ).transpose(1, 0, 2, 3)

        # 1. Primary Loss: Ensures each expert head is trained well.
        all_head_losses = []
        for i, h in enumerate(self.config.horizons):
            v_t_head = all_v_t_preds[i, :, :h, :]
            target_v_t = u_t[:, :h, :]
            # Mean over batch, horizon, and action dim
            head_loss = jnp.mean(jnp.square(v_t_head - target_v_t))
            all_head_losses.append(head_loss)

        individual_loss = jnp.sum(jnp.stack(all_head_losses))

        # 2. Auxiliary Loss: Trains the gating network
        # (B*H, S_suffix, 1)
        gate_logits_batched = self.gate_out_proj(suffix_out)
        # (B*H, max_H, 1)
        gate_logits_padded = gate_logits_batched[:, action_start_index: action_start_index + max_horizon, :]
        # (B, max_H, H)
        # gate_logits = einops.rearrange(gate_logits_padded, "(b h) s 1 -> b s h", b=batch_size, h=num_horizons)
        gate_logits_reshaped = gate_logits_padded.reshape(batch_size, num_horizons, max_horizon, 1)
        gate_logits = jnp.transpose(gate_logits_reshaped, (0, 2, 1, 3)).squeeze(-1)

        # Create mask for softmax
        # (max_H, H)
        valid_heads_mask = jnp.arange(max_horizon)[:, None] < horizons_arr[None, :]
        # (B, max_H, H) - broadcast batch dim
        valid_heads_mask = jnp.broadcast_to(valid_heads_mask, gate_logits.shape)

        masked_gate_logits = jnp.where(valid_heads_mask, gate_logits, jnp.finfo(gate_logits.dtype).min)
        gate_weights = nnx.softmax(masked_gate_logits, axis=-1)

        # Combine predictions using the dynamic weights
        # all_v_t_preds: (H, B, max_H, D) -> (B, H, max_H, D)
        all_v_t_preds_permuted = einops.rearrange(all_v_t_preds, "h b s d -> b h s d")
        # gate_weights: (B, max_H, H) -> (B, H, max_H, 1)
        gate_weights_expanded = jnp.transpose(gate_weights, (0, 2, 1))[:, :, :, None]

        # (B, H, max_H, D) * (B, H, max_H, 1) -> sum over H -> (B, max_H, D)
        v_t_combined = jnp.sum(all_v_t_preds_permuted * gate_weights_expanded, axis=1)

        auxiliary_loss = jnp.mean(jnp.square(v_t_combined - u_t))  # Mean over B, H, D

        # 3. Balance Loss: Encourage the gate layer to output weights flexibly
        loss_components = []
        boundaries = sorted(list(set([0] + self.config.horizons)))
        for i in range(len(boundaries) - 1):
            start_step, end_step = boundaries[i], boundaries[i + 1]
            active_expert_indices = [idx for idx, h in enumerate(self.config.horizons) if h > start_step]

            if len(active_expert_indices) > 1:
                # (B, S_segment, H_total)
                segment_gate_weights = gate_weights[:, start_step:end_step, :]
                # (B, S_segment, H_active)
                active_expert_weights = segment_gate_weights[:, :, jnp.array(active_expert_indices)]
                # (H_active,)
                avg_expert_prob_in_segment = jnp.mean(active_expert_weights, axis=(0, 1))
                segment_loss = self.cv_squared(avg_expert_prob_in_segment)
                loss_components.append(segment_loss)

        load_balancing_loss = jnp.mean(jnp.stack(loss_components)) if loss_components else 0.0

        total_loss = (
                individual_loss +
                self.config.aux_weight * auxiliary_loss +
                self.config.balance_weight * load_balancing_loss
        )

        return total_loss

    @override
    def sample_actions(
        self,
        rng: at.KeyArrayLike,
        observation: _model.Observation,
        *,
        num_steps: int | at.Int[at.Array, ""] = 10,
        noise: at.Float[at.Array, "b ah ad"] | None = None,
    ) -> _model.Actions:
        """
        Samples actions using the gated fusion mechanism during denoising.
        """
        observation = _model.preprocess_observation(None, observation, train=False)
        dt = -1.0 / num_steps
        batch_size = observation.state.shape[0]
        max_horizon = self.action_horizon
        num_horizons = len(self.config.horizons)
        horizons_arr = jnp.array(self.config.horizons)

        noise = jax.random.normal(rng, (batch_size, max_horizon, self.action_dim))

        # First fill KV cache with a forward pass of the prefix
        prefix_tokens, prefix_mask, prefix_ar_mask = self.embed_prefix(observation)
        prefix_attn_mask = make_attn_mask(prefix_mask, prefix_ar_mask)
        positions = jnp.cumsum(prefix_mask, axis=1) - 1
        (_, prefix_out), kv_cache = self.PaliGemma.llm(
            [prefix_tokens, None],
            mask=prefix_attn_mask,
            positions=positions
        )

        # Prepare static batched inputs (these don't change in the loop)
        batched_prefix_mask = jnp.repeat(prefix_mask, num_horizons, axis=0)
        batched_kv_cache = jax.tree_map(
            lambda x: jnp.repeat(x, num_horizons, axis=1),
            kv_cache
        )
        batched_state = jnp.repeat(observation.state, num_horizons, axis=0)

        # Create static masks for padding actions in the loop
        # (H, max_H)
        steps_arr = jnp.arange(max_horizon)
        action_pad_mask_per_h = steps_arr[None, :] < horizons_arr[:, None]
        # (B*H, max_H)
        batched_action_pad_mask = jnp.broadcast_to(
            action_pad_mask_per_h[None, :, :],
            (batch_size, num_horizons, max_horizon)
        )
        batched_action_pad_mask = einops.rearrange(batched_action_pad_mask, "b h s -> (b h) s")
        # batched_action_pad_mask = einops.repeat(action_pad_mask_per_h, "h s -> (b h) s", b=batch_size)
        # (H, max_H, 1)
        action_mask_for_padding_x_t = (steps_arr[None, :, None] < horizons_arr[:, None, None])

        # Create static mask for gate softmax
        # (max_H, H)
        valid_heads_mask = steps_arr[:, None] < horizons_arr[None, :]
        # (B, max_H, H) - for broadcasting
        valid_heads_mask = valid_heads_mask[None, :, :]
        
        action_start_index = 0 if self.pi05 else 1  # pi0.5 has no state token

        prefix_len = prefix_mask.shape[1]

        def step_fn(carry):
            x_t, time = carry

            # --- Prepare Batched Inputs for this step ---
            expanded_time = jnp.broadcast_to(time, (batch_size * num_horizons,))

            # Pad x_t for each horizon
            # (1, B, max_H, D)
            x_t_expanded = x_t[None, ...]
            # (H, B, max_H, D)
            padded_x_t_batched = jnp.where(action_mask_for_padding_x_t, x_t_expanded, 0.0)
            # (B*H, max_H, D)
            batched_x_t = padded_x_t_batched.transpose(1, 0, 2, 3)
            batched_x_t = einops.rearrange(batched_x_t, "b h s d -> (b h) s d")

            # --- Run Batched Suffix Pass ---
            suffix_tokens, suffix_pad_masks, suffix_ar_mask, adarms_cond = self.embed_suffix(
                batched_state, batched_x_t, expanded_time, action_pad_mask=batched_action_pad_mask
            )

            pad_masks = jnp.concatenate([batched_prefix_mask, suffix_pad_masks], axis=1)
            ar_masks = jnp.concatenate([prefix_ar_mask, suffix_ar_mask], axis=0)
            full_att_2d_masks = make_attn_mask(pad_masks, ar_masks)

            suffix_len = suffix_tokens.shape[1]
            suffix_position_ids = jnp.arange(prefix_len, prefix_len + suffix_len)[None, :]
            suffix_att_2d_masks = full_att_2d_masks[:, -suffix_len:, :]

            b = suffix_att_2d_masks.shape[0]
            suffix_position_ids = jnp.broadcast_to(suffix_position_ids, (b, suffix_len))
            
            adarms = [None, adarms_cond] if self.pi05 else [None, None]
            (_, suffix_out), _ = self.PaliGemma.llm(
                [None, suffix_tokens],
                mask=suffix_att_2d_masks,
                positions=suffix_position_ids,
                kv_cache=batched_kv_cache,
                adarms_cond=adarms,
            )

            # --- Gating and Fusion ---
            # (B*H, S_suffix, 1)
            gate_logits_batched = self.gate_out_proj(suffix_out)
            # (B*H, max_H, 1)
            gate_logits_padded = gate_logits_batched[:, action_start_index: action_start_index + max_horizon, :]
            # (B, max_H, H)
            gate_logits_reshaped = gate_logits_padded.reshape(batch_size, num_horizons, max_horizon, 1)
            gate_logits = jnp.transpose(gate_logits_reshaped, (0, 2, 1, 3)).squeeze(-1)
            masked_gate_logits = jnp.where(valid_heads_mask, gate_logits, jnp.finfo(gate_logits.dtype).min)
            gate_weights = nnx.softmax(masked_gate_logits, axis=-1)

            # Get all predictions
            # (B*H, S_suffix, D_action)
            v_t_batched = self.action_out_proj(suffix_out)
            # (B*H, max_H, D_action)
            v_t_actions_padded = v_t_batched[:, action_start_index: action_start_index + max_horizon, :]
            # (B, H, max_H, D_action)
            all_v_t_preds = v_t_actions_padded.reshape(batch_size, num_horizons, max_horizon, -1)

            # Combine predictions
            # gate_weights: (B, max_H, H) -> (B, H, max_H, 1)
            gate_weights_expanded = jnp.transpose(gate_weights, (0, 2, 1))[:, :, :, None]

            # (B, H, max_H, D) * (B, H, max_H, 1) -> sum over H -> (B, max_H, D)
            v_t = jnp.sum(all_v_t_preds * gate_weights_expanded, axis=1)

            # --- Euler Step ---
            x_t_new = x_t + dt * v_t
            time_new = time + dt

            return (x_t_new, time_new)

        def cond_fn(carry):
            x_t, time = carry
            # robust to floating-point error
            return time >= -dt / 2

        x_0, _ = jax.lax.while_loop(cond_fn, step_fn, (noise, 1.0))
        return x_0