3DLLM-Mem: Long-Term Spatial-Temporal Memory for Embodied 3D Large Language Model

We appreciate your valuable comments, which are crucial for improving our paper! Thank you for finding our "3D benchmark is novel" and acknowledging that "the idea of working memory and episodic memory is interesting."

W1.1 & Q1: Efficiency comparison with other methods.

Thank you for this insightful question. We've benchmarked the efficiency on the medium set of our embodied tasks and calculated the average inference time per question, and the results show that our method achieves a massive improvement in success rate with only a marginal increase in inference time.

In-domain Tasks

In-the-wild Tasks

The main computational cost for all methods comes from the 7B LLM backbone. The slight difference in overhead arises from the memory processing step:

• Most Recent Memory: No processing is needed; the memory is fed directly to the LLM.
• Retrieval-Augmented: Overhead comes from calculating similarity scores against the memory bank.
• Our Method: Overhead comes from a single cross-attention layer.

The complexity of our fusion step of the query ($N$ x $d$) and past memories key and values which are ($T$ x $N$ x $d$) is $O(TN^2d)$. Since the memory bank size $T$ is bounded by our task definition (Line 142), the $N$ and $d$ shares the same across all methods, this operation is highly efficient. Furthermore, its performance could be easily accelerated using standard optimizations like FlashAttention.

W1.2: Moreover, if the thinking process is too long, the environment will be changed a lot, and the past observations and interactions might be useless.

Thank you for raising this point. We respectfully suggest that for long-horizon reasoning tasks, past observations often remain highly relevant, even after significant time has passed.

For example, in the task shown in Figure 1, the agent must remember the size of the first teddy bear it found to correctly select a matching gift box much later in the episode. If it were to forget this initial information, the task would become unsolvable. The "thinking process" our model performs is precisely what allows it to connect these temporally distant but causally linked events. We are, of course, open to further discussion on scenarios where this might not hold.

W2: Limited to static environment, while cannot handle the dynamic environment.

Thank you for this concern. We'd like to clarify that our framework is indeed designed for and evaluated in dynamic environments. This is a core aspect of our work, mentioned in Lines 3, 7, and 64.

The agent is constantly interacting with and changing the environment by executing actions (as seen in Figures 1 and 3). Our Memory Update mechanism (Section 3, Lines 202-209) is specifically designed to handle this. When the agent alters a part of the scene, the memory entry for that location is updated with a new observation, ensuring the model's knowledge of the world stays current.

W3: Ablation of 3DLLM-Mem effectiveness with base model, the memory setting, the training data.

Our short answer is that the effectiveness comes from our novel dynamic memory management and fusion module.

To isolate this contribution, our experimental setup uses the same base model (LLaVA-3D) and the same training dataset for all methods. We then compare several memory management techniques: "Everything in Context," "Most Recent Memory," "Retrieval-Augmented Memory," and our proposed 3DLLM-Mem.

As demonstrated in Table 2, our method significantly outperforms all other baselines under these identical conditions. This shows that the performance gain is attributable to our memory fusion design, not the base model or the data. We provide further targeted ablations on the components of our fusion module in Table 3 to reinforce this conclusion.

W4: Failure Cases

For the failure cases, we have identified three common categories:

1. Perception Misalignment: The alignment of the 3D perception encoder with the Large Language Model (LLM) remains a bottleneck in this domain, a challenge we also observed. When a room contains noisy depth data, incomplete 3D meshes, or small/occluded objects, our model can fail to perceive the target objects correctly, leading to incorrect subsequent actions.
2. Premature Task Completion: We append a Task Complete token to the end of every training sequence to signal task termination. For complex tasks requiring multiple steps (e.g., collecting four instances of the same object from different rooms), the agent can sometimes lose track of its progress and incorrectly output the Task Complete token before the goal is fully achieved. We believe employing an explicit "counting" module could help the model track this information, which is an interesting direction for future work.
3. Cascading Errors: We were pleasantly surprised to find that our model can sometimes recover from an incorrect action mid-task and still successfully complete its goal. However, a more common scenario is that an initial mistake (e.g., picking up the wrong object) leads to a cascade of subsequent incorrect actions. We believe that incorporating reinforcement learning with sparse rewards to penalize incorrect actions could improve model robustness and is another promising area for future research.

Q2: How to adopt this model into scenes with more dynamics and interaction?

Please see Weakness 2. Our setting is already dynamic, and the model is designed to handle interactions that change the environment.

Q3: What is the contribution of each component in your framework?

Please see Weakness 3. The primary contribution comes from our memory fusion module. We also provide a detailed ablation study in Table 3 that breaks down the effectiveness of its specific design choices.

Limitation.

We mentioned limitations in our conclusion (Lines 327-330), but we agree this section can be expanded. We will add a more detailed discussion to the appendix in our final version, including the failure cases listed in our response to W4. Thank you for the suggestion.

We sincerely appreciate your thoughtful comments and hope our responses have addressed your concerns. Please let us know if you have any further questions.

Best, Authors

After reading the rebuttal, I think most parts of my concerns have been addressed. However, I still think this work has not provided sufficient evidences for the dynamic scenes, and I hope the authors can address it in the future.

In summary, I keep the positive score.

We'd like to express our sincere gratitude for your thorough review of our paper. We're thrilled that you found our work "tackles a significant and underexplored challenge," that our "benchmark is a great contribution," and that our "experimental evaluation is convincing." We greatly appreciate your suggestions, which are crucial for improving the quality of our paper.

W1: Contextalizing our work w.r.t prior work.

Thank you for pointing out “Memorizing Transformers” (ICLR ’22). While our initial literature survey focused on embodied AI, we agree this is a relevant work and will cite and discuss it in our revised manuscript. Its approach highlights key differences that underscore the novelty of our method:

1. Memery Fusion vs. Retrieval: “Memorizing Transformers” utilizes a k-NN lookup to retrieve top (key, value) pairs from memory. This is a form of retrieval, similar in principle to our "Retrieval-Augmented Memory" baseline. Our proposed 3DLLM-Mem method is fundamentally different; it uses fusion. While retrieval selects a subset of the "best" memories based on a fixed similarity score, our fusion module attends to all memories simultaneously. It uses a parameterized cross-attention mechanism (Line 196) to learn a new, task-relevant representation by selectively aggregating information from the entire memory bank. This allows the model to learn what past information is most important for the current step, a capability our experiments show is highly effective.
2. Dynamic vs. Static Memory: “Memorizing Transformers” stores a long-term context (like a text document) as a static KV cache. In our embodied setting, the memory is dynamic. As the agent interacts with the environment, the world state changes (Line 202, Memory Update). Our model is designed to handle this dynamism by continuously updating its memory bank and fusing these updated memories at each step to inform its actions, as shown in Figure 3.
3. Different Problem Settings: As you acknowledged, the settings are distinct. “Memorizing Transformers” focuses on extending the effective context window for language models on text-based tasks. Our work, as stated in our abstract, introduces "a novel dynamic memory management and fusion model for embodied spatial-temporal reasoning and actions."

Moreover, these significant differences in application settings fundamentally shape model design choices—for example, the approach to extending a language transformer’s context window via KV cache versus enabling spatial-temporal memory for dynamic 3D embodied tasks.

W2: Presentation: What is the task definition and what are the model’s inputs and outputs? The 'what' is missed from paper's Introductory section.

Thank you for highlighting this. We apologize if the introduction was not perfectly clear. We will definitely revise the introduction to be more explicit.

To clarify here:

• The task is for an embodied agent, given a high-level language instruction (e.g., "prepare the most suitable gift box for the teddy bear."), to navigate a 3D environment and perform a sequence of actions to complete the goal. We provide a visual "teaser" of this in Figure 1.
• The model's input at each step is the current 3D visual observation (working memory) and a bank of past 3D observations (episodic memory). Please also refer to paper's Sec 1, Line 71.
• The model's output is a symbolic action, such as PickUp(teddy bear).

We will incorporate this clear problem definition into the introductory section of our revised paper.

W3: Presentation: A preliminary section briefly describing the LLAVA-3D and clear problem setting and task definition section before Section 2.

Thank you for this excellent structural suggestion. We'd like to clarify that we do introduce LLaVA-3D as a preliminary in Section 3.1 (Lines 165-176).

However, we agree with your point about the overall structure. Our original intention was to present our benchmark (which includes the task settings) before the model. Based on your feedback, we see that introducing the high-level problem setting and task definition before Section 2 would improve the paper's flow. We will make this structural change in our final version.

W4 & Q2: Why is the proposed method superior to the Retrieval-Augmented Memory baseline?

You are exactly right that "the key difference is in how the memory fusion is performed." As we elaborated in our response to W1, the core distinction is Fusion vs. Retrieval.

To summarize: The Retrieval-Augmented baseline selects the top-K "best" memories based on a static similarity metric. This has two potential drawbacks: the similarity metric may not be optimal, and the top-K memories might not contain all the necessary context.

In contrast, our fusion-based method attends over all memories and learns a dynamic, task-dependent way to combine their information. It isn't limited to a few retrieved memories; it can synthesize information from the entire history. This allows it to capture more complex, long-range dependencies and makes it more robust, as it's not dependent on a single, potentially flawed retrieval step.

Reponse to Nit:

Thank you so much for your careful reading! We will happily make these changes for clarity in the revised version:

• We'll change “reason over long-term memory in 3D environments” to “reason over extended temporal spans using long-term memory.”
• We'll change the column name "Long-term Memory" in Table 1 to something less ambiguous, such as "Evaluates Memory," to avoid confusion.

Q1: Clarify how precisely memery update is implemented/executed?

We are happy to further explain the memory update process (Lines 202-209).

First, as noted in Line 142, the number of distinct locations (e.g., rooms) in any task is finite, so the memory bank doesn't grow indefinitely. The memory bank itself is a structured collection of feature vectors, where each vector corresponds to a specific, previously visited location.

When the agent revisits a location and performs an action that alters the scene (e.g., picking up an object), we re-compute the 3D feature representation for that location. We then update the corresponding entry in the memory bank with this new feature vector. Implementation wise, it's as simple as Bank(entry)=new_vector. This is what's described in Lines 205-207: "If the corresponding environment already exists in the memory bank and has been modified by the agent, the memory entry is updated accordingly." This process ensures the agent's memory of the world state remains current.

Limitation:

Thank you for this valuable suggestion. We did have to shorten our limitation section due to space constraints, but we will add a more comprehensive discussion to the Appendix, including real-world evaluation and a detailed breakdown of failure cases.

As a preview, here are the three common failure categories we've identified:

1. Perception Misalignment: The alignment of the 3D perception encoder with the Large Language Model (LLM) remains a bottleneck in this domain, a challenge we also observed. When a room contains noisy depth data, incomplete 3D meshes, or small/occluded objects, our model can fail to perceive the target objects correctly, leading to incorrect subsequent actions.
2. Premature Task Completion: We append a Task Complete token to the end of every training sequence to signal task termination. For complex tasks requiring multiple steps (e.g., collecting four instances of the same object from different rooms), the agent can sometimes lose track of its progress and incorrectly output the Task Complete token before the goal is fully achieved. We believe employing an explicit "counting" module could help the model track this information, which is an interesting direction for future work.
3. Cascading Errors: We were pleasantly surprised to find that our model can sometimes recover from an incorrect action mid-task and still successfully complete its goal. However, a more common scenario is that an initial mistake (e.g., picking up the wrong object) leads to a cascade of subsequent incorrect actions. We believe that incorporating reinforcement learning with sparse rewards to penalize incorrect actions could improve model robustness and is another promising area for future research.

We sincerely appreciate your thoughtful comments and hope our responses have addressed your concerns. Please let us know if you have any further questions.

Best, Authors

Thank you for your careful reading and detailed feedback. We're glad you found our proposed memory fusion mechanism to be "novel and well-justified", and we appreciate you highlighting our model's "strong results and generalization to unseen environments".

We provide detailed explanations for the points of clarity you raised below and will update our paper accordingly to ensure these concepts are easier to understand. We're also happy to engage further during the discussion period to clarify any remaining questions.

W1 & Q1: What the difference is between the in-domain vs in-the-wild tasks? What are “instances” and “unseen memory context”?

Thank you for this question. Let's first clarify the terms to facilitate the explanation.

• Instance: An "instance" is simply a single data point or evaluation episode.
• Environment: An "environment" refers to a specific 3D scene, such as a particular house layout with its rooms and objects. As illustrated in Figure 1, the agent navigated from the bedrooms to the living room, and then to the kitchen—all of which constitute the environment.
• Unseen Memory Context: This means the agent is operating in an environment that it has never encountered during training.

With these definitions, the distinction between our evaluation splits (please also refer to paper's Lines 144-148) is as follows:

1. From an Environment/Memory Perspective:

• In-domain tasks take place in environments that may have been seen during training. Therefore, the agent's episodic memory bank might contain memories from a familiar scene. However, the agent's current observation (its working memory) is always from a novel, unseen state within that environment.
• In-the-wild tasks, as noted in Line 147, feature "entirely unseen memory context." This means the agent operates in completely new environments that were not part of the training data.

2. From a Task Difference Perspective:

• In-domain tasks follow the same distribution and types of tasks as those in the training set. For example, as illustrated in Figure 2, one such task is "Find an ideal container to store all the cookie dough."
• In-the-wild tasks, as noted in Line 148, introduce novel challenges that are more complex and qualitatively different from those seen during training. Examples of these challenging tasks are provided in Figure 2; corresponding to the in-domain example above, for instance, is "Collect all cookie dough and then rearrange it on the living room table in descending order of size."

A brief note on design: We designed the in-domain set this way to ensure the in-the-wild set remained a reasonably sized and robust evaluation pool containing our most valuable, entirely novel instances. The in-domain tasks are still challenging, as shown by the experimental results, because the agent must still act based on its immediate, unseen working memory.

W2: The definitions of the environment, working memory, and episodic memory are unclear.

We apologize for the confusion regarding the terms "working memory" and "episodic memory," which are inspired by concepts in neuroscience. Your understanding that "working memory stores a single timestep" is exactly right! It represents the agent's perception of its current 3D environment at a specific moment.

We believe the confusion may have stemmed from the sentence: "When the agent moves to a new environment, the previous working memory is transferred to the episodic memory bank." To be more precise: at each timestep $t$, the agent's current observation is its working memory. When the agent proceeds to timestep $t$+1, the observation from the previous step $t$ becomes a part of its past experience. This past observation is then incorporated into the episodic memory bank. This process is visualized in Figure 3(b), where the working memory tokens $f_t^Q$ are fused with the past memory bank features.

To further emphasize this distinction, our paper defines their structures with mathematical notation:

• Working Memory (Line 182) has a shape of $R^{(N \times d)}$, representing features from one time step.
• Episodic Memory (Line 184) has a shape of $R^{(T \times N \times d)}$, representing features from $T$ past time steps.

We hope this resolves the confusion, and we will revise the text in the paper to make this relationship clearer

W3: Details of training are missing.

We appreciate you raising this, and we would like to clarify where these details are located in the paper. We describe our training setup in Section 4.1, under the paragraph Implementation Details (Lines 218-224). This section covers the training computation, architecture, and loss function.

For more granular details, we mention in Line 224 that a full list of hyperparameters is available in Appendix D (Lines 691-699). There, we specify the learning rate, optimizer, weight decay, and learning rate scheduler used for our experiments. We will also open-source our code after the review period. Thank you for your understanding.

W4.1: While 3DLLM-Mem has much higher absolute performance on in-the-wild tasks, the relative drop in performance between in-domain and in-the-wild is only slightly lower than other method.

Thank you for this insightful observation. To better address your point, we've created a new table using the scores from Table 2, adding columns for absolute and relative performance drops.

As illustrated above, our method demonstrate a -5.5 absolute and a -14.6% relative performance drop which is significantly smaller than the >30% relative drop seen with other methods. We believe this demonstrates a substantial improvement in generalization, not just a "slight" one.

Furthermore, we'd argue that relative drop is a more equitable metric here. It's inherently more challenging to maintain a high level of performance (like our 37.6% -> 32.1%) on difficult, unseen tasks than it is to maintain a low level of performance (like 21.1% -> 13.7%). Our model's ability to retain most of its high performance highlights its robustness.

W4.2: More insight into this would strengthen this contribution, such as explaining which components of the method lead to this improvement.

Thank you for recognizing our method's effectiveness. We attribute this strong generalization capability directly to our novel memory fusion module (Section 3.2), which excels for two key reasons:

1. Relevant Initialization: The fusion process begins with query tokens initialized from the agent's current working memory. This is crucial because it primes the attention mechanism to selectively retrieve and focus on past memories that are most relevant to the agent's immediate surroundings and task. As shown in our ablation study (Table 3), initializing from other sources leads to inferior performance.
2. Holistic Fusion Process: Unlike simple retrieval methods that might only fetch the single "best" memory, our fusion module attends to all past memories simultaneously. It learns to weigh and synthesize them into a single, compact, and context-rich representation. This allows the model to capture complex, long-term dependencies and make more informed decisions, which is especially critical for the novel and challenging tasks in the in-the-wild set.

Q2: How are the following terms defined: instance, environment, unseen memory context?

Please see Weakness 1 and 2.

Q3: What is the relationship between the working and episodic memory?

Please see Weakness 2. In short, working memory is the agent's current observation, which becomes part of the episodic (past) memory in the next timestep.

Q4: Do you have insight into why 3DLLM-Mem demonstrates strong generalization capabilities? Is this simply due to better overall performance? Or are there specific components of the method that enable it to maintain performance compared to other approaches?

Please see Weakness 4.2. The performance improvement is from our deisgned fusion memory module as introduced in paper's section 3.2.

We sincerely appreciate your thoughtful comments and hope our responses have addressed your concerns. Please let us know if you have any further questions.

Best, Authors

We sincerely thank you for your constructive feedback and for describing 3DMem-Bench as a "novel and fine-grained" benchmark and 3DLLM-Mem "effectively improves the performance of the model in complex scenarios". We greatly appreciate your suggestions which are crucial in improving the quality of our paper.

W1: How to handle scenarios when memory bank growing indefinitely?

Thank you for raising this valuable concern. In the current benchmark configuration (Sec.2, L.142), an embodied task involves at most 10 rooms. Yet we have constructed a challenging benchmark as all models' performance is still low. Therefore, in the context of 3DMem-Bench, the memory bank remains small enough to avoid computational inefficiency.

Nevertheless, we agree that scalability is an essential consideration for broader real-world applications. We are happy to discuss the following promising approaches to manage memory growth:

1. Sliding-window memory (FIFO). Maintain only the most recent N memories and discard older ones. This is similar to the “Most-Recent Memory” baseline in our paper but now applies to the memory bank. This method keeps both latency and memory O(N).
2. Memory consolidation / compression. The cross attention complexity of the query ($N$ x $d$) and past memories key and values which are ($T$ x $N$ x $d$) is $O(TN^2d)$. We can: 1) down-sample token length $N$. 2) project to a lower dimension $d$ (e.g., LoRA-style adapters), or 3) cluster similar memories and retain only their representative centroids. While promising, a careful study of the resulting accuracy–efficiency trade-off would be an excellent topic for future work.
3. Retrieval-augmented subsets. Index memories with a light-weight key–value store (FAISS, ScaNN) and then apply memory fusion only to the top-K retrieved memories. This reduces the effective $T$ and is orthogonal to the above techniques.

We believe all these approaches can be promising future works but beyond the scope of the current paper.

W2: The tasks in a simulator may not fully capture the physical world.

Thank you for this feedback. We share this concern and directly investigated the generalization and transferability of our model within the 3DMem-Bench by splitting tasks into "in-domain" and "in-the-wild" sets. Our results show that 3DLLM-Mem demonstrates strong generalization capabilities on the unseen "in-the-wild" tasks.

Moreover, we believe our approach has high potential for real-world transferability. As illustrated in paper's Figure 1 and 3, the model outputs high-level, symbolic actions (e.g., Navigate(room=kitchen), PickUp(bottle)). The low-level motion primitives are executed by modules that are already deployed and tested on real robots. Consequently, the learned policy primarily needs to decide what to do, not how to do it, which substantially narrows the reality gap often encountered in methods that learn low-level continuous control policies.

Q1: Failure Cases.

Thank you for this question. We'll clarify the failure cases for both the data generation process and our trained model.

Regarding the task generation process, as you noted, we discuss these limitations in Appendix F.1. The primary failure modes involve Gemini occasionally generating trajectories that violate constraints related to: 1) the agent’s final location, 2) the specific object referenced, and 3) the correctness of pick-up and put-down actions. The single most common failure is violating the "hold only one object at a time" constraint, as this requires strong reasoning and a consistent memory of the agent's current inventory.

For the failure cases of our trained 3DLLM-Mem model, we have identified three common categories:

1. Perception Misalignment: The alignment of the 3D perception encoder with the Large Language Model (LLM) remains a bottleneck in this domain, a challenge we also observed. When a room contains noisy depth data, incomplete 3D meshes, or small/occluded objects, our model can fail to perceive the target objects correctly, leading to incorrect subsequent actions.
2. Premature Task Completion: We append a Task Complete token to the end of every training sequence to signal task termination. For complex tasks requiring multiple steps (e.g., collecting four instances of the same object from different rooms), the agent can sometimes lose track of its progress and incorrectly output the Task Complete token before the goal is fully achieved. We believe employing an explicit "counting" module could help the model track this information, which is an interesting direction for future work.
3. Cascading Errors: We were pleasantly surprised to find that our model can sometimes recover from an incorrect action mid-task and still successfully complete its goal. However, a more common scenario is that an initial mistake (e.g., picking up the wrong object) leads to a cascade of subsequent incorrect actions. We believe that incorporating reinforcement learning with sparse rewards to penalize incorrect actions could improve model robustness and is another promising area for future research.

We sincerely appreciate your thoughtful comments and hope our responses have addressed your concerns. Please let us know if you have any further questions.

Best, Authors