# Notes for future SAE activation capture (GR00T-N1.7-LIBERO)

> Status: **not implemented yet.** This file just records the hook points found
> while wiring up the LIBERO smoke tests, so a later change can add SAE
> activation collection with minimal surface area. Nothing here changes model
> behaviour.

## Where actions are produced (the model call chain)

1. `gr00t/eval/rollout_policy.py :: run_rollout_gymnasium_policy`
   calls `policy.get_action(observations)` once per action chunk.
   - For the smoke tests the `policy` is a `PolicyClient`
     (`gr00t/policy/server_client.py`) talking over ZMQ to a server, so the
     *actual* model lives in the server process. The
     `examples/LIBERO/smoke_tests/_libero_rollout_worker.py` already wraps
     `PolicyClient.get_action` to record the returned action chunks
     (`actions.npy`); that wrapper is the cheapest place to record
     *input/output* of the policy from the client side, but it cannot see
     internal activations.

2. Server side: `gr00t/eval/run_gr00t_server.py` builds a
   `Gr00tPolicy` (`gr00t/policy/gr00t_policy.py`), optionally wrapped in
   `Gr00tSimPolicyWrapper`. The endpoint `get_action` →
   `Gr00tPolicy._get_action()` (≈ line 371):
   - builds `collated_inputs` from the observation (images + proprio state +
     language instruction), casts to bf16, then:
   - `model_pred = self.model.get_action(**collated_inputs)` ← **the forward pass**
   - `normalized_action = model_pred["action_pred"]`, then
     `self.processor.decode_action(...)` → physical-unit action dict returned.

3. The model: `gr00t/model/gr00t_n1d7/gr00t_n1d7.py :: GR00T_N1_7` (registered
   `Gr00tN1d7`). `GR00T_N1_7.get_action(inputs)` (≈ line 589):
   ```
   backbone_inputs, action_inputs = self.prepare_input(inputs)
   backbone_outputs = self.backbone(backbone_inputs)          # VLM
   action_outputs   = self.action_head.get_action(backbone_outputs, action_inputs, options)
   return action_outputs                                       # {"action_pred": ...}
   ```

## Likely hook points / module names

Top-level model attributes (`policy.model` on the server, an
`nn.Module` of class `Gr00tN1d7` / `GR00T_N1_7`):

| Attribute | Class / file | What it computes |
|---|---|---|
| `model.backbone` | `Qwen3Backbone` (`gr00t/model/modules/qwen3_backbone.py`) | VLM over images + text instruction. Wraps HF `Qwen3VLForConditionalGeneration` at `model.backbone.model`; transformer decoder layers at `model.backbone.model.language_model.layers[...]` (truncated to `select_layer`); vision tower at `model.backbone.model.visual` (name may vary with the HF Qwen3-VL version). `forward()` runs with `output_hidden_states=True`, takes `hidden_states[-1]`, and returns `BatchFeature({"backbone_features": <Tensor [B, seq, hidden]>, "backbone_attention_mask": ...})`. |
| `model.action_head` | `Gr00tN1d7ActionHead` (`gr00t/model/gr00t_n1d7/gr00t_n1d7.py`, ≈ line 38) | Flow-matching / diffusion action decoder. Inner net `model.action_head.model` is a `DiT` or `AlternateVLDiT` (`gr00t/model/modules/dit.py`, `gr00t/model/modules/flowmatching_modules.py`). Has `get_action(backbone_output, action_input, options)` and `get_action_with_features(...)` (≈ line 312) which is the natural place to also surface intermediate features. |
| `model.collator` | `Gr00tN1d7DataCollator` (`gr00t/model/gr00t_n1d7/processing_gr00t_n1d7.py`) | builds VLM batch inputs from `vlm_content`. |

**Recommended primary capture site:** the output of `model.backbone` — i.e.
the tensor stored under key `backbone_features` returned by
`Qwen3Backbone.forward` (last LLM hidden state, shape `[B, seq_len, hidden]`).
This is the single fused representation of *image + instruction + (implicit)
proprioception context* that conditions action generation, so it's the most
informative single activation to feed an SAE. Secondary sites: per-layer LLM
hidden states (`outputs.hidden_states[k]` inside `Qwen3Backbone.forward`),
and the DiT block activations inside `model.action_head.model`.

**Cheapest mechanism:** register `torch.nn.Module.register_forward_hook` on
`policy.model.backbone` (and/or specific `...language_model.layers[k]`) right
after the policy is constructed in `run_gr00t_server.py`, before
`server.run()`. Hooks can append detached CPU/float16 tensors to a buffer that
is flushed to disk per get_action call. No edits to model code required.
Alternative: subclass `Gr00tPolicy` and override `_get_action` to also stash
`self.model`'s intermediates.

## Where observation / instruction / proprioception enter the model

- **LIBERO env → observation dict**: `gr00t/eval/sim/LIBERO/libero_env.py ::
  LiberoEnv._process_observation` produces:
  - `video.image` (agentview 256×256×3), `video.wrist_image` (eye-in-hand),
  - `state.x/y/z/roll/pitch/yaw` (EEF pose from `robot0_eef_pos` + axis-angle
    of `robot0_eef_quat`), `state.gripper` (2-dim `robot0_gripper_qpos`),
  - `annotation.human.action.task_description` (the **language instruction**,
    a fixed string per task).
- These are batched/temporally-stacked by `MultiStepWrapper`
  (`gr00t/eval/sim/wrapper/multistep_wrapper.py`) and then, in
  `Gr00tPolicy._get_action`, turned into `VLAStepData` via
  `self._to_vla_step_data(obs)` and run through `self.processor(messages)`
  (`Gr00tN1d7Processor`, `gr00t/model/gr00t_n1d7/processing_gr00t_n1d7.py`):
  - images + instruction → `vlm_content` → tokenized/pixel-processed by the
    Qwen3-VL processor → `model.backbone`.
  - proprio `state.*` → `action_inputs` → `model.action_head` (conditioning).
- The smoke-test worker can optionally perturb the image observations
  *before* `policy.get_action` (`--obs-noise-std`); a future activation-capture
  run could pair clean vs. perturbed activations from the same seeds.

## Suggested next step for activation capture

1. Add an opt-in flag to `run_gr00t_server.py` (e.g. `--capture-activations
   DIR` / `--capture-layers backbone,llm.20`) that, after building the policy,
   registers forward hooks on the chosen modules and writes one
   `.npz`/`.safetensors` per `get_action` call (keyed by an episode/step id
   passed through `options`).
2. Have `_libero_rollout_worker.py` thread a stable `(scenario_id, episode,
   step)` tag into `policy.get_action(obs, options=...)` so captured
   activations can be joined back to `actions.npy` and the saved video frames.
3. Train the SAE on `backbone_features` first (one SAE), then expand to
   per-layer LLM hidden states if needed.