`internal_state` not used in rollout

[SHAO-Jiaqi757]

In the paper, Figure 3 describes the RL pipeline and rollout for MEM1:

“Figure 3: (Top): the RL pipeline used to train MEM1. (Bottom left): The evolution of context in MEM1–old <IS>, <query>, <info> are cleared as new states enter the context. The mechanism is used in the rollout.”

However, in the current implementation:

• internal_state only occurs in Mem1/inference/data_pipelines.py, within Mem1Pipeline.run_llm_loop for inference only.
  https://github.com/MIT-MI/MEM1/blob/main/Mem1/inference/data_pipelines.py#L109-L143

• internal_state is not referenced in the rollout (Mem1/train/rollout/llm_agent/generation_think.py).
  https://github.com/MIT-MI/MEM1/blob/main/Mem1/train/rollout/llm_agent/generation_think.py#L264-L290

[BobbyZhouZijian]

Hi. Thank you for your question. During rollout, the internal state is stored as f"t{step}" in the reconstruction_list variable.

https://github.com/MIT-MI/MEM1/blob/main/Mem1/train/rollout/llm_agent/generation_think.py#L377-L381

[SHAO-Jiaqi757]

Hi. Thanks for your reply!

If I understood correctly:

1. During Rollout

• The model generates responses in a truncated context.
• _extract_think_and_response extracts the <think> section.
• This section is stored as reconstruction_list[i][f't{step}'] (you mention as internal_state).

2. During Training

• flatten_trajectory marks all t{step} entries with kind=1.
• make_attention_mask applies special attention rules for tokens with kind=1.

Purpose of reconstruction_list

1. To save all generated content (without truncation) during rollout.
2. To reconstruct the complete trajectory (including all observations) during training.

About context saving

The “context saving” is because: During inference, previous turns’ observations are truncated via the kept_lengths mechanism,
—not because of internal state preservation. internal state didn't "replace" the previous observations as described in the paper.

MEM1/Mem1/train/rollout/llm_agent/tensor_helper.py

Lines 39 to 48 in 4aea4c5

	def mask_using_kept_lengths(self, attn_mask: torch.Tensor, kept_lengths: torch.Tensor, initial_token_lengths: List[int]) -> torch.Tensor:
	for i in range(len(kept_lengths)):
	# find the last non-zero index in the i-th row of attn_mask
	non_zero_indices = (attn_mask[i] == 1).nonzero()
	last_non_zero_index = non_zero_indices[-1]
	first_non_zero_index = non_zero_indices[0]
	attn_mask[i, :last_non_zero_index - kept_lengths[i] + 1] = 0
	attn_mask[i, first_non_zero_index:first_non_zero_index + initial_token_lengths[i]] = 1
	return attn_mask

[lujiarui-iie]

And in training rollout phase, when generating the think for step j, the think for step j-1 was masked, which is inconsistent with the 2D Attention Mask in Figure 3.