OmniJARVIS: Unified Vision-Language-Action Tokenization Enables Open-World Instruction Following Agents

Q5: Results comparison with Jarvis-1 and Voyager.

Thank you for your inquiry. Jarvis-1 and Voyager were not included in Table 2 due to their testing frameworks and controller:

1. Testing Settings: Jarvis-1 and Voyager operate under few-shot settings, utilizing multiple inference cycles and an explicit textual memory module for life-long learning. In contrast, our experiments are conducted in a zero-shot setting, focusing on each agent’s ability to follow instructions and generalize without prior exposure.

2. Controller Differences: Voyager employs a scripted action (offered by MineFlyer) and accesses privileged environmental information (voxel and lidar data), which provides it with capabilities not shared by the baseline agents in our study, which primarily use policy-based and visual perceptive controllers.

These differences make direct comparisons between the models unfair and uninformative. We will clarify this in the updated manuscript.

Q6: Using Codebook to derive interpretable skills.

Thanks for the suggestions. As mentioned in the General Response, OmniJARVIS primarily aims at offering a more compact representation of skills in VLA compared to counterparts like RT-2 that employ language annotations, which can be expensive to obtain. We commit to exploring the interpretability of our behavior codebook as part of future work.

W1, Q1, and L2: Impact of interaction dataset and model architecture.

Thank you for your insightful comments. To clarify the individual contributions and combined impact of the new dataset and architectural methods, we conducted specific ablation studies detailed in our paper.

Dataset Contributions: OmniJARVIS utilizes two key datasets: a text QA dataset based on Minecraft knowledge and the Contractor dataset of human gameplay in Minecraft. The text QA dataset enhances OmniJARVIS’s understanding of Minecraft knowledge and high-level planning for complex tasks. The Contractor dataset, on the other hand, improves mastery of basic Minecraft skills such as “mine oak log.” We explored the impact of these datasets on different task types through new ablation experiments presented in Table 2 of the supplementary PDF. Results indicate that the QA dataset significantly influences OmniJARVIS’s performance on long-horizon programmatic tasks, while the Contractor dataset facilitates success in short-horizon atom tasks.

Architectural Innovations: To assess the contributions of architectural innovations, we also conducted experiments on the effects of the vision encoder (Fuyu, LLaVA, isolated Captioner), behavior tokenizer (VQ-VAE, FSQ, Language), and various scales and architectures of the language model (Gemma-2B, LLaMA-7B, LLaMA-13B). These are detailed in Tables 4 and 6, and Figure 4 of the paper.

Our findings demonstrate that both the datasets and architectural enhancements significantly contribute to OmniJARVIS’s ability to perform a wide range of tasks, from simple atom tasks to complex programmatic challenges. The synergy between the diverse datasets and innovative architectural elements is key to the robust performance of OmniJARVIS across different task domains.

W2: Comprehensive ablations on FSQ and dataset.

Thank you for your comments. We’ve conducted several ablation studies to validate the individual contributions of our architecture:

1. FSQ Ablations: We’ve added new experiments on FSQ codebook size and code context length, detailed in the PDF-Table-1 and General Response.

2. Dataset Ablations: We add new dataset ablation in Table 2 of the supplementary PDF to investigate the effectiveness of the QA and interaction dataset. Detailed in Table 5 of the original paper, these experiments evaluate the impact of different dataset segments on OmniJARVIS’s performance.

3. Behavior Tokenizer Ablations: Experiments on different tokenizer types (language, VQ-VAE, FSQ) are added to Table 3 of supplementary PDF, highlighting their effects on behavior generation.

These studies rigorously assess the influence of each component on the model’s effectiveness.

Q2 and L1: Ablations on multimodal LLM.

Thank you for the suggestion to include more comprehensive ablations and comparisons with multimodal LLMs. We have conducted experiments using various visual encoders, including FUYU (patch encoder), LLAVA (ViT), and an independent vision captioning model (ShareCaptioner+). The results of these experiments are detailed in Table 4 of the main paper.

Additionally, we are currently finetuning the OmniJARVIS model using the QFormer-based BLIP2-OPT-2.7B architecture. This variant of OmniJARVIS has not yet been finished; we will update the manuscript with the results once they are available.

Q3: Details on behavior tokenizer ablation in Table 6.

Thanks for your comments. This table showcases the success rates of OmniJARVIS in completing programmatic tasks using different behavior tokenizers.

Regarding your specific question about language goals, they indeed include memory and caption text derived from the meta-information in the contractor data. These are converted into language goals that serve as decoder policy prompts for the language-conditioned STEVE-I model. For example, a language goal could be "chop down the tree to mine oak logs."

A note on the VQ-VAE tokenizer: during its training, we encountered a posterior collapse where only one code from the VQ codebook was utilized, leading to the failure of the VQ-GROOT training. The FSQ-GROOT, in contrast, represents the final setting used for OmniJARVIS.

All models in this comparison, including LLaVA-7B as the base model, used the same synthetic memory, thought, and caption data to ensure fairness in our evaluations. We continue to expand the scale of the ablation experiments (more tasks and test repetitions) and have placed the latest results in supplementary PDF Table 3.

Q4: More evaluation times in Table 6.

In Table 6, each task was evaluated 10 times. Due to the time-intensive nature of Programmatic Tasks, which require 6000-12000 timestamps to complete, this number of evaluations was deemed a practical balance between comprehensiveness and feasibility. To expand the scale of our experiments, we selected 1-2 representative tasks from each group of Programmatic Tasks and tested each 20 times. The results of this expanded testing are presented in Table 3 of the supplementary PDF, offering a larger sample size and further validation of our findings.

Due to word limit restrictions, we have added more responses to the official comments, and we hope you can see them.

W6 and Q1: OmniJARVIS on other environments.

Thank you for your suggestions. We have indeed begun extending OmniJARVIS to other environments, starting with Atari Montezuma’s Revenge game, where it achieved a score of 3600. This initial success illustrates the model’s potential for generalization. Details on adapting OmniJARVIS to different environments are further elaborated in the General Response. We are committed to ongoing efforts to expand its applicability across various domains, demonstrating its versatility and broader utility.

W7: Add limitations.

Thanks for your comments. We discuss the limitations of OmniJARVIS in the General Response, which will be inserted in the updated manuscript.

W8 and Q6: More analysis on agent failure modes.

Thank you for your inquiry regarding the failure modes of OmniJARVIS and its behaviors, particularly in challenging tasks such as obtaining a diamond. The primary failure modes can be categorized into the following:

1. Planning Errors: During programmatic tasks, OmniJARVIS occasionally produces incorrect thoughts or plans. For instance, it might attempt to mine stone without a wooden pickaxe. These errors stem from inaccuracies within the QA dataset and hallucinations inherent in the language model.

2. Action Execution Failures: The FSQ GROOT decoder (our low-level control policy) sometimes fails to translate the behavior code into actions properly, particularly with less frequently occurring fsq codes due to data imbalance in the interaction dataset.

3. Hallucinations in perception: Errors in visual recognition by the Llava model can lead to incorrect scene interpretations, such as mistaking a stone for iron ore. These hallucinations hinder the agent’s reasoning and decision-making processes.

In the specific task of obtaining a diamond, these failure modes manifest when OmniJARVIS incorrectly plans or executes sequences due to the above issues. However, OmniJARVIS has a high success rate on finishing such tasks and correctly sequences tasks from mining basic materials to crafting necessary tools and finally obtaining the diamond, demonstrating effective integration of vision, language, and action within one unified model.

Q3: why not re-using language tokens?

Thank you for your question about integrating behavior tokens into the LLM vocabulary. We employ different strategies based on the tokenizer used:

1. Reserved Tokens: When possible, we use reserved tokens from the language tokenizer for behavior tokens to maintain semantic integrity. This happens with the tokenizer of Fuyu VLM.

2. Reusing Infrequently Used Tokens: In the absence of suitable reserved tokens, we repurpose infrequently used tokens, as seen with the Llama tokenizer. This happens with the LLaVA tokenizer. We apologize for not making this clear and have updated the manuscript.

Furthermore, our FSQ design compresses extensive codebook sizes into up to 35 tokens (when the FSQ setting is 8+8+8+6+5), optimizing vocabulary use without semantic loss. These approaches balance semantic preservation with efficient vocabulary management.

Q4: Confusion on training dataset tokens.

Thank you for your attention to detail. The 1T tokens were indeed a typo, the correct calculation should be 1B behavior and language tokens in total. We apologize for the confusion and will correct this in the revised manuscript.

Q5: How OmniJARVIS works without instruction?

Thank you for your question. In the absence of instructions, OmniJARVIS effectively models a dataset with gameplay video only -- the videos are segmented into trunks with a default length of 128, and then all trunks are converted to codes by the behavior tokenizer. OmniJARVIS is tasked to predict these codes based on the initial visual observations of the corresponding segments. The goal of this ablation is to verify the necessity of including plans, thoughts, and other means of instruction (the "language" modality) in the VLA modeling of OmniJARVIS.

Minor 2: checkmark in training settings

Apologies for the confusion. We use a checkmark to indicate when the model's training data includes specific information. In Q5, the first row without an instruction represents unconditional OmniJARVIS, which is unable to follow human language instructions. With richer synthetic data, including thought, memory, and caption data, OmniJARVIS can better follow instructions, leading to improved learning with lower loss. This demonstrates the effectiveness of synthetic data. We will follow your suggestion to replace the symbols.

W1 and Q2: Details on encoder and decoder of behavior tokenizer.

Sorry for the confusion. As shown in Fig. 2, the Behavior Tokenizer consists of three parts: Encoder, Decoder, and FSQ quantizer. We perform quantization based on the original GROOT, and the Encoder and Decoder use a design consistent with GROOT. Specifically, the encoder includes a Convolutional Neural Network (CNN) to extract spatial information from image states $s_{1:T}$ and a non-causal transformer to capture temporal information from videos. The CNN is EfficientNet-B0 networks and the non-causal transformer is built on the code from minGPT-2 removing the causal mask. The decoder consists of 4 identical blocks, where each block contains a Transformer-XL block and a Flamingo gated-attention dense layer to condition the decoder on FSQ code embeddings. More details can be found in GROOT paper Section 4.1 and Section 4.2. The fsq quantizer consists of 5 levels of [8,8,8,6,5] with a codebook size of 15360. The training hyperparameters of FSQ-GROOT are as follows: Parallel Strategy: ddp, Accumulate Gradient Batches: 8, Batch size: 2, Precision: bf16, Image size 224*224, Encoder blocks: 8, Decoder blocks: 8, Hidden dimension: 1024, Chunk size: 128, Attention memory size: 256, Optimizer: AdamW, Weight decay: 0.001, Learning rate: 1.81e-05, Warmup steps: 2k.

W2: Ablation experiments on behavior tokenizer (codebook size and context length).

Thanks for your comments. We conduct an in-depth investigation of behavior tokenizers with varying codebook sizes, utilizing recommended sets of FSQ levels to approximate specified codebook dimensions in the supplementary PDF Table-1. Codebook Usage is quantified as the proportion of codewords utilized at least once when encoding the validation datasets. Reconstruction FSD is measured by the FSD scores derived from the MineCLIP encoder, processing 1,000 different demonstration videos through the FSQ-GROOT and subsequent rollout in a randomly generated environment. Additionally, we measure Resampling FSD, which is the FSD score obtained when the environment rollout is conditioned on representations sampled from the codebook. Finally, we assess the average rewards for the task “collect wood” using OmniJARVIS across varying codebook sizes. Our findings indicate that increases in codebook size correlate with enhanced average rewards and reduced FSD scores, suggesting a scalable performance in OmniJARVIS with larger codebooks.

We also conduct an ablation of experiment on different context length of behavior token in PDF Table 1. We select the codebook size e14 as the default setting and explore different context lengths of 128, and 64. The performances under context lengths of 64 and 128 are similar. Our insight is, that although shorter context lengths may offer better behavior granularity, they also require more frequent behavior code switching during inference, resulting in a slower speed of OmniJARVIS.

W2: Behavior tokens are only conditioned on observation sequence.

The reason we use observation sequence as tokenizer input is as follows: we inherit the self-supervised learning scheme from GROOT to train the behavior tokenizer, and the learning objective is to estimate missing actions from the observation-only sequence. Including actions in the input will break this learning scheme.

W3 and Minor1: Evaluation details of programmatic tasks.

Sorry for the confusion. The number behind the task category is the number of tasks in this group. The evaluated programmatic tasks are taken from DEPS, and the detailed task description and task settings are shown in the supplementary PDF Table 5. All tasks are evaluated over 20 times from an empty inventory under the open-ended generated environment in Minecraft. We will add this section to the updated manuscript.

W4: Ablation experiments on the training dataset.

Thank you for your comments. To validate the effectiveness of including the QA dataset, we conducted ablation studies (detailed in Supplementary PDF Table 2). These studies revealed that while OmniJARVIS performs similarly on Atom Tasks with or without the QA dataset, there is a significant performance drop in Programmatic Tasks when the QA dataset is omitted. This drop underscores the necessity of the QA dataset for enhancing reasoning and planning capabilities essential for Programmatic Tasks.

We will further clarify this aspect in the updated manuscript, ensuring a comprehensive understanding of the QA dataset’s contribution to OmniJARVIS’s enhanced performance compared to baselines.

W5: Performance drop of OmniJARVIS in Table 1.

Thank you for your comments. We apologize for the confusion caused by an error in Table 1; the success rate for OmniJARVIS on the stone task should be 15.8, not 5.8 as mistakenly listed. This typo has skewed the comparison.

Furthermore, it’s important to note that GROOT relies on manually selected reference videos as instructions, which can significantly influence its performance. In contrast, OmniJARVIS autonomously generates behavior tokens for the policy decoder to execute, providing a more consistent and scalable approach. We further tested more atom tasks to compare GROOT and OmniJARVIS in the following table. Our tests show that OmniJARVIS performs comparably to GROOT across most Atom tasks, demonstrating its robustness not just in programmatic tasks but also in short-horizon atom tasks.

Due to word limit restrictions, we have added more responses to the official comments, and we hope you can see them.

W1: difference between GROOT and OmniJARVIS

As detailed in the General Response, our core contribution is to propose a novel Vision-Language-Action (VLA) model architecture termed OmniJARVIS, as shown in Figure 1 of the supplementary PDF, aiming at resolving the issues including inference efficiency and language and annotation bottleneck.

GROOT is limited to short-horizon tasks (atom tasks) and cannot accept language instructions, as demonstrated in Tables 1 and 2 in the original paper. OmniJARVIS, on the other hand, uses the FSQ-GROOT decoder as a de-tokenizer, enabling it to generate environment-acceptable actions from discrete FSQ codes that are produced by the VLM. This allows OmniJARVIS to handle complex, long-horizon tasks (programmatic tasks) with advanced reasoning and planning abilities. Please let us know if you still have further concerns about the novelty and we're more than happy to assist!

W2: Descriptions for OmniJARVIS training.

Sorry for the confusion. We utilized the SFTTrainer class from the TRL library by Hugging Face to train the VLM model. The learning rate was set at 1.4e-5, and a cosine learning rate scheduler was employed. The weight decay parameter was set to 0 with a warm-up ratio of 0.03. Training took place on 8 A800 GPUs with FSDP, with a batch size of 2 and gradient accumulation steps of 4 using bf16 precision. The training lasted for one epoch on our generated dataset.

W3: typos

Sorry for the typos. We've revised the manuscript accordingly.

W4: How to handle longer trajectories (>128).

Thank you for your question. OmniJARVIS is specifically designed to model trajectories that exceed this length, which is a limitation in models like GROOT which can only handle fixed 128-frame segments for instruction-following.

Given a long trajectory(>128), we first slice it into multiple segments of 128 frames each. These segments are then encoded into action tokens, $D^{beh}$, and integrated into an interaction dataset formulated as sequences of ${D^{instruct}, D^{mem}, D^{obs}, D^{th}, D^{beh}, D^{obs}, D^{th}, D^{beh}, ...}$. OmniJARVIS, as an autoregressive transformer model, supports a maximum token length of 8k, allowing it to model sequences that contain multiple $D^{behavior}$ tokens, effectively handling long sequences beyond 128 frames.

Additionally, during training, we employ a sliding window technique to manage sequences that exceed 8k tokens in VLM context length. In inference, OmniJARVIS outputs a new behavior token every 128 frames to address long-horizon tasks such as programmatic and creative tasks.

Q1 and Q2: Programmatic tasks are included in OmniJARVIS training datasets?

OmniJARVIS leverages two datasets for training: a Question-answering dataset enriched with Minecraft knowledge and the Contractor dataset of human gameplay trajectories. While the instructions in the Contractor dataset don’t exclusively pertain to programmatic tasks, they support OmniJARVIS in generalizing from learned knowledge to generate new plans and behaviors in varied instructional contexts. For the comparisons in Table 2, all models, including the baselines, were trained using the same Minecraft dataset to ensure fairness in evaluating generalization capabilities across different tasks.

Additionally, the baseline models in Table 2 were not retrained or finetuned with the data constructed specifically for OmniJARVIS. However, we conducted further tests, shown in Table 3, finetuning open-source models on a Minecraft-specific Question-Answering dataset derived via ChatGPT self-instruct. This helped evaluate the impact of injecting specialized Minecraft knowledge on model performance, where the finetuned models showed improved understanding, aligning their capabilities closer to those of ChatGPT.

Q3: Discussion on GROOT performance on Atom and Programmatic tasks.

Atom tasks are short-horizon tasks that can usually be completed within 600 timestamps (30s). These tasks primarily assess the agent’s mastery of basic skills in Minecraft, such as “mine oak log” and “kill sheep.” In contrast, programmatic tasks are long-horizon tasks, often requiring more than 6000 timestamps (5 min) to complete. These tasks demand complex reasoning and planning. For example, the programmatic task “obtain diamond” involves multiple intermediate steps, such as mining oak logs, crafting a wooden pickaxe, and acquiring stone and iron ore.

GROOT and STEVE-I models are trained on the Minecraft Contractor data for instruction-following tasks of fewer than 128 frames. As a result, they possess certain atom skill proficiency but lack complex planning and reasoning capabilities as opposed to VLA models like OmniJARVIS. Consequently, they perform well on simple atom tasks but struggle with programmatic tasks that require sophisticated planning and reasoning.

Q4: Explanation of Figure 5.

Figure 5 aims to visualize the behavior produced by our FSQ-GROOT decoder (the low-level policy adopted by OmniJARVIS) when conditioned on certain behavior codes. On the left is a screenshot of the input reference video to be encoded into behavior codes, and on the right is a screenshot of the FSQ-GROOT decoder policy's rollout in a new environment. The behavioral similarity between left and right screenshots verifies that the code produced by the behavior tokenizer can be consistently decoded into the same behavior as the input by decoder policy, even in novel environments.

W1: Novelty concerns.

Thank you for your comments. As detailed in the General response, our core contribution is to propose a novel Vision-Language-Action (VLA) model architecture termed OmniJARVIS, as shown in Figure 1 of the supplementary PDF, aiming at resolving the issues including inference efficiency and language and annotation bottleneck. Please let us know if you still have further concerns about the novelty and we're more than happy to assist!

W2: Training/Inference efficiency analysis and comparison.

Thank you for your valuable feedback. We’ve added comparisons of inference speed, memory usage, and model parameters to Table 4 in the supplementary PDF:

Inference Efficiency Comparison: The inference efficiency experiments are conducted on the same workstation with one RTX3090 GPU.

1. Inference Speed: Based on Minecraft’s rendering speed of approximately 25 fps, we found that the inference speed of models using a native action tokenizer (direct action output), such as GROOT and VPT, inversely correlates with model size. The same architecture’s VPT model shows decreasing inference speeds with increasing parameters: 1x, 2x, and 3x.
2. Memory Usage: Models based on language tokenizers, which require in-context learning facilitated by ChatGPT, exhibited the lowest inference speeds due to their complex computational requirements.
3. OmniJARVIS Performance: Thanks to its self-supervised behavior tokenizer and hierarchical architecture, OmniJARVIS achieves faster inference speeds (16.62 fps) even with a larger parameter count (7.2 billion) compared to native tokenizer models like VPT-3x (0.5 billion parameters at 12.34 fps).

Training Efficiency: OmniJARVIS enhances training efficiency by encoding N=128 frames of trajectory data into k=5 behavior tokens. This compression allows the model to handle data more effectively compared to models like RT-2 and OpenVLA, which predict actions directly. This results in a significant reduction in training data volume and corresponding training time.

These quantitative results have been included in the revised paper to substantiate our claims regarding the efficiency of OmniJARVIS compared to other approaches.

Q: cost for synergized dataset.

Due to the relatively certain generation of Thought, Memory, and Instruction, we used gpt-3.5-turbo to synthesize this data, and processing the Contractor dataset cost around $400.

L: Add a separate "Limitations" section and Generalization of OmniJARVIS.

Thank you for your suggestion. We have added a dedicated "Limitations" section in the General Response to address the scope of validation for OmniJARVIS. Initially, the model was validated solely within the Minecraft environment. Recognizing the need for broader applicability, we have also conducted preliminary generalization experiments in Atari Montezuma’s Revenge environment. These results, discussed in the General Response, demonstrate OmniJARVIS’s potential for adaptation to different scenarios. We will provide further details on these aspects in the updated manuscript to ensure a comprehensive understanding of the model’s limitations and capabilities.