How to support multi round chat training?

[xuexidi]

Thank you very much for your open-source work!
I found that UI-R1 only supports single image and single round dialogue training. If I want to support multi image and multi round chat training, how do I modify the code?
My multi round chat data sequence is roughly as follows:
[
user_text1
user_image1
assistant_response1
user_image2
assistant_response2
user_image3
assistant_response3
user_image4
...
]

[lll6gg]

You could modify the function make_conversation_from_json in line 264 in src/ui_r1/src/open_r1/grpo_json_action_coord.py to read a multi-round and multi-image chat from a json file.

[xuexidi]

You could modify the function make_conversation_from_json in line 264 in src/ui_r1/src/open_r1/grpo_json_action_coord.py to read a multi-round and multi-image chat from a json file.

But I carefully read the Qwen2VLGRPOTraining class from open_r1.trainer import Qwen2VLGRPOTraining and found that the def _generate_1and_stcore_completions function in the source code of the class involves dialogue processing. I found that the function only reads one image from a certain input:

for x in inputs:
           if "image" in x:
               img = x["image"]
           elif "image_path" in x and x["image_path"] is not None:
               img = PIL.Image.open(x["image_path"])

Doesn't this mean that even if I construct multiple rounds of conversation data by modifying make_comversation_from_json, the Qwen2VLGRPOTrainingclass actually only processes one image in each conversation, right?

The function details of _generate_and_stcore_completions are as follows:

    def _generate_and_score_completions(self, inputs: dict[str, Union[torch.Tensor, Any]], model) -> dict[str, Union[torch.Tensor, Any]]:
        device = self.accelerator.device
        prompts = [x["prompt"] for x in inputs]
        prompts_text = [maybe_apply_chat_template(example, self.processing_class)["prompt"] for example in inputs]
        # Handle both pre-loaded images and image paths
        images = []
        for x in inputs:
            if "image" in x:
                img = x["image"]
            elif "image_path" in x and x["image_path"] is not None:
                img = PIL.Image.open(x["image_path"])

            try:
                # Ensure minimum dimensions of 28 pixels
                w, h = img.size
                if w < 28 or h < 28:
                    # Calculate new dimensions maintaining aspect ratio
                    if w < h:
                        new_w = 28
                        new_h = int(h * (28/w))
                    else:
                        new_h = 28
                        new_w = int(w * (28/h))
                    img = img.resize((new_w, new_h), PIL.Image.Resampling.LANCZOS)
            
                images.append(img)
            except:
                pass

        if len(images) > 0:
            prompt_inputs = self.processing_class(
                text=prompts_text,
                images=images,
                return_tensors="pt",
                padding=True,
                padding_side="left",
                add_special_tokens=False,
            )
        else:
            prompt_inputs = self.processing_class(
                text=prompts_text,
                return_tensors="pt",
                padding=True,
                padding_side="left",
                add_special_tokens=False,
            )
        prompt_inputs = super()._prepare_inputs(prompt_inputs)

        # resize output coordinate due to the image resize
        origin_height = images[0].size[1]
        origin_width = images[0].size[0]

        # option 1
        # resized_height, resized_width = smart_resize(origin_height, origin_width, max_pixels=self.processing_class.image_processor.max_pixels)
        # option 2
        resized_height = prompt_inputs['image_grid_thw'][0][1] * self.processing_class.image_processor.patch_size
        resized_width = prompt_inputs['image_grid_thw'][0][2] * self.processing_class.image_processor.patch_size
        
        scale_x = origin_width / resized_width
        scale_y = origin_height / resized_height
        scales = [scale_x,scale_y]
        prompt_ids, prompt_mask = prompt_inputs["input_ids"], prompt_inputs["attention_mask"]
        if len(images) > 0:
            pixel_values = prompt_inputs["pixel_values"]
            image_grid_thw = prompt_inputs["image_grid_thw"]
        else:
            pixel_values = None
            image_grid_thw = None

        
        # if self.max_prompt_length is not None:
        #     prompt_ids = prompt_ids[:, -self.max_prompt_length :]
        #     prompt_inputs["input_ids"] = prompt_ids
        #     prompt_mask = prompt_mask[:, -self.max_prompt_length :]
        #     prompt_inputs["attention_mask"] = prompt_mask

        # Generate completions
        with unwrap_model_for_generation(model, self.accelerator) as unwrapped_model:
            prompt_completion_ids = unwrapped_model.generate(
                **prompt_inputs, 
                generation_config=self.generation_config
            )

            prompt_length = prompt_ids.size(1)
            prompt_ids = prompt_completion_ids[:, :prompt_length]
            completion_ids = prompt_completion_ids[:, prompt_length:]
            # No need to repeat prompt_mask as we're not duplicating prompts during generation
        ......

[lll6gg]

I got it. You could solve it by:

1. changing image_path to image_paths (List) in make_conversation_from_json .
2. reading image_paths and appending them into images in _generate_and_score_completions. The inputs will consist of several images. The key code is:

        if len(images) > 0:
            prompt_inputs = self.processing_class(
                text=prompts_text,
                images=images,
                return_tensors="pt",
                padding=True,
                padding_side="left",
                add_special_tokens=False,
            )

For detail, you could imitate the Qwen2.5-VL and change the code.

[xuexidi]

I got it. You could solve it by:

1. changing image_path to image_paths (List) in make_conversation_from_json .
2. reading image_paths and appending them into images in _generate_and_score_completions. The inputs will consist of several images. The key code is:

        if len(images) > 0:
            prompt_inputs = self.processing_class(
                text=prompts_text,
                images=images,
                return_tensors="pt",
                padding=True,
                padding_side="left",
                add_special_tokens=False,
            )

For detail, you could imitate the Qwen2.5-VL and change the code.

I found that the implementation logic in the function '_generate_and_stcore_completions' already looks like the following code:

# Handle both pre-loaded images and image paths
        images = []
        for x in inputs:
            if "image" in x:
                img = x["image"]
            elif "image_path" in x and x["image_path"] is not None:
                img = PIL.Image.open(x["image_path"])

            try:
                # Ensure minimum dimensions of 28 pixels
                w, h = img.size
                if w < 28 or h < 28:
                    # Calculate new dimensions maintaining aspect ratio
                    if w < h:
                        new_w = 28
                        new_h = int(h * (28/w))
                    else:
                        new_h = 28
                        new_w = int(w * (28/h))
                    img = img.resize((new_w, new_h), PIL.Image.Resampling.LANCZOS)
            
                images.append(img)
            except:
                pass

        if len(images) > 0:
            prompt_inputs = self.processing_class(
                text=prompts_text,
                images=images,
                return_tensors="pt",
                padding=True,
                padding_side="left",
                add_special_tokens=False,
            )

However, this processing logic is to append all the images of the Q&A in a batch to the same images: list. If I append multiple images from the same conversation to the images: list, wouldn't all the images from the conversation be mixed together and messy? (Moreover, for my multi round dialogue dataset mentioned above, there are strict order requirements for the input user_dempt, assistant_response, and image.)

Did I not understand something correctly？

Thank you for your patient reply!

[xuexidi]

The images: list represents the single QA but not a batch.

Huh?
"def _generate_and_score_completions(self, inputs: dict[str, Union[torch.Tensor, Any]], model) -> dict[str, Union[torch.Tensor, Any]]" Isn't the 'inputs' of this function corresponding to the inputsof a batch?
Isn't the following code taking the prompt of Q for all QAs in the current batch?

prompts = [x["prompt"] for x in inputs]
prompts_text = [maybe_apply_chat_template(example, self.processing_class)["prompt"] for example in inputs]

    def _generate_and_score_completions(self, inputs: dict[str, Union[torch.Tensor, Any]], model) -> dict[str, Union[torch.Tensor, Any]]:
        device = self.accelerator.device
        prompts = [x["prompt"] for x in inputs]
        prompts_text = [maybe_apply_chat_template(example, self.processing_class)["prompt"] for example in inputs]
        # Handle both pre-loaded images and image paths
        images = []
        for x in inputs:
            if "image" in x:
                img = x["image"]
            elif "image_path" in x and x["image_path"] is not None:
                img = PIL.Image.open(x["image_path"])

[lll6gg]

Sorry i made a mistake, I think it's better to maintain another variable image_idxs to record the chat ids of the images.

[xuexidi]

Sorry i made a mistake, I think it's better to maintain another variable image_idxs to record the chat ids of the images.

Thank you, do you know how to specify the chat ID of the corresponding image for the following code?
I am not familiar with the usage of 'processing_class'. Where can I find detailed instructions？

if len(images) > 0:
    prompt_inputs = self.processing_class(
        text=prompts_text,
        images=images,
        return_tensors="pt",
        padding=True,
        padding_side="left",
        add_special_tokens=False,
    )

[lll6gg]

batch item : (prompt,images) may be better. The images list and prompt are one-to-one. Exampe:

def _generate_and_score_completions(self, inputs: dict[str, Union[torch.Tensor, Any]], model) -> dict[str, Union[torch.Tensor, Any]]:
    device = self.accelerator.device
    prompts_text = []
    images_list = []
    for example in inputs:
        prompt = maybe_apply_chat_template(example, self.processing_class)["prompt"]
        prompt_images = []
        for message in example.get("prompt", []):
            for content in message.get("content", []):
                if content["type"] == "image":
                    if "image" in content and content["image"] is not None:
                        img = content["image"]
                    elif "image_path" in content and content["image_path"] is not None:
                        img = PIL.Image.open(content["image_path"])
                    else:
                        continue
                    w, h = img.size
                    if w < 28 or h < 28:
                        if w < h:
                            new_w = 28
                            new_h = int(h * (28 / w))
                        else:
                            new_h = 28
                            new_w = int(w * (28 / h))
                        img = img.resize((new_w, new_h), PIL.Image.Resampling.LANCZOS)
                    prompt_images.append(img)
        prompts_text.append(prompt)
        images_list.append(prompt_images)
    prompt_inputs = self.processing_class(
        text=prompts_text,
        images=images_list,
        return_tensors="pt",
        padding=True,
        padding_side="left",
        add_special_tokens=False,
    )
    prompt_inputs = super()._prepare_inputs(prompt_inputs)
    with unwrap_model_for_generation(model, self.accelerator) as unwrapped_model:
        prompt_completion_ids = unwrapped_model.generate(
            **prompt_inputs, 
            generation_config=self.generation_config
        )
        prompt_length = prompt_inputs['input_ids'].size(1)
        prompt_ids = prompt_completion_ids[:, :prompt_length]
        completion_ids = prompt_completion_ids[:, prompt_length:]
    return {
        "prompt_ids": prompt_ids,
        "completion_ids": completion_ids
    }

You could double check and have a try.

[xuexidi]

batch item : (prompt,images) may be better. The images list and prompt are one-to-one. Exampe:

def _generate_and_score_completions(self, inputs: dict[str, Union[torch.Tensor, Any]], model) -> dict[str, Union[torch.Tensor, Any]]:
    device = self.accelerator.device
    prompts_text = []
    images_list = []
    for example in inputs:
        prompt = maybe_apply_chat_template(example, self.processing_class)["prompt"]
        prompt_images = []
        for message in example.get("prompt", []):
            for content in message.get("content", []):
                if content["type"] == "image":
                    if "image" in content and content["image"] is not None:
                        img = content["image"]
                    elif "image_path" in content and content["image_path"] is not None:
                        img = PIL.Image.open(content["image_path"])
                    else:
                        continue
                    w, h = img.size
                    if w < 28 or h < 28:
                        if w < h:
                            new_w = 28
                            new_h = int(h * (28 / w))
                        else:
                            new_h = 28
                            new_w = int(w * (28 / h))
                        img = img.resize((new_w, new_h), PIL.Image.Resampling.LANCZOS)
                    prompt_images.append(img)
        prompts_text.append(prompt)
        images_list.append(prompt_images)
    prompt_inputs = self.processing_class(
        text=prompts_text,
        images=images_list,
        return_tensors="pt",
        padding=True,
        padding_side="left",
        add_special_tokens=False,
    )
    prompt_inputs = super()._prepare_inputs(prompt_inputs)
    with unwrap_model_for_generation(model, self.accelerator) as unwrapped_model:
        prompt_completion_ids = unwrapped_model.generate(
            **prompt_inputs, 
            generation_config=self.generation_config
        )
        prompt_length = prompt_inputs['input_ids'].size(1)
        prompt_ids = prompt_completion_ids[:, :prompt_length]
        completion_ids = prompt_completion_ids[:, prompt_length:]
    return {
        "prompt_ids": prompt_ids,
        "completion_ids": completion_ids
    }

You could double check and have a try.

Thank you very much! I'll take a try~

By the way，I found the SYSTEM_PROMPT in grpo_json_action_coord.py were never used, so when I test my trained Qwen2.5-VL-3B, should the model system prompt be left as default (you are a helpful assistant.)?

SYSTEM_PROMPT = (
    "A conversation between User and Assistant. The user asks a question, and the Assistant solves it. The assistant "
    "first thinks about the reasoning process in the mind and then provides the user with the answer. The reasoning "
    "process and answer are enclosed within <think> </think> and <answer> </answer> tags, respectively, i.e., "
    "<think> reasoning process here </think><answer> answer here </answer>"
)

[lll6gg]

Yes I didn't use it. You could add {"role": "system", "content": SYSTEM_PROMPT} to make_conversation_from_json. It doesn't matter so much.

[xuexidi]

Yes I didn't use it. You could add {"role": "system", "content": SYSTEM_PROMPT} to make_conversation_from_json. It doesn't matter so much.

OK，I see，Thanks!

[xuexidi]

batch item : (prompt,images) may be better. The images list and prompt are one-to-one. Exampe:

def _generate_and_score_completions(self, inputs: dict[str, Union[torch.Tensor, Any]], model) -> dict[str, Union[torch.Tensor, Any]]:
    device = self.accelerator.device
    prompts_text = []
    images_list = []
    for example in inputs:
        prompt = maybe_apply_chat_template(example, self.processing_class)["prompt"]
        prompt_images = []
        for message in example.get("prompt", []):
            for content in message.get("content", []):
                if content["type"] == "image":
                    if "image" in content and content["image"] is not None:
                        img = content["image"]
                    elif "image_path" in content and content["image_path"] is not None:
                        img = PIL.Image.open(content["image_path"])
                    else:
                        continue
                    w, h = img.size
                    if w < 28 or h < 28:
                        if w < h:
                            new_w = 28
                            new_h = int(h * (28 / w))
                        else:
                            new_h = 28
                            new_w = int(w * (28 / h))
                        img = img.resize((new_w, new_h), PIL.Image.Resampling.LANCZOS)
                    prompt_images.append(img)
        prompts_text.append(prompt)
        images_list.append(prompt_images)
    prompt_inputs = self.processing_class(
        text=prompts_text,
        images=images_list,
        return_tensors="pt",
        padding=True,
        padding_side="left",
        add_special_tokens=False,
    )
    prompt_inputs = super()._prepare_inputs(prompt_inputs)
    with unwrap_model_for_generation(model, self.accelerator) as unwrapped_model:
        prompt_completion_ids = unwrapped_model.generate(
            **prompt_inputs, 
            generation_config=self.generation_config
        )
        prompt_length = prompt_inputs['input_ids'].size(1)
        prompt_ids = prompt_completion_ids[:, :prompt_length]
        completion_ids = prompt_completion_ids[:, prompt_length:]
    return {
        "prompt_ids": prompt_ids,
        "completion_ids": completion_ids
    }

You could double check and have a try.

I just looked at the work of VLM-R1, which supports multiple image input methods（ https://github.com/om-ai-lab/VLM-R1/blob/c04f284fa2c0f847c9b05f3fd22e6695fb1ec9fe/src/open-r1-multimodal/src/open_r1/trainer/grpo_trainer.py#L528 ）The writing here is to put all the images in the images: list, without distinguishing different images according to the prompt. This seems to cause multiple images corresponding to a prompt to not be paired and aligned during training. Is this a bug in VLM-R1's work?

def _generate_and_score_completions(self, inputs: dict[str, Union[torch.Tensor, Any]], model) -> dict[str, Union[torch.Tensor, Any]]:
        device = self.accelerator.device
        prompts = [x["prompt"] for x in inputs]
        prompts_text = self.vlm_module.prepare_prompt(self.processing_class, inputs)
        # Handle both pre-loaded images and image paths
        images = []
        for x in inputs:
            if "image" in x:
                imgs = self._get_key_from_inputs(x, "image")
            elif "image_path" in x and x["image_path"] is not None:
                imgs = [PIL.Image.open(p) for p in self._get_key_from_inputs(x, "image_path")]
            else:
                imgs = []

            for img in imgs:
                try:
                    # Ensure minimum dimensions of 28 pixels
                    w, h = img.size
                    if w < 28 or h < 28:
                    # Calculate new dimensions maintaining aspect ratio
                        if w < h:
                            new_w = 28
                            new_h = int(h * (28/w))
                        else:
                            new_h = 28
                            new_w = int(w * (28/h))
                    img = img.resize((new_w, new_h), PIL.Image.Resampling.LANCZOS)
                except:
                    pass
                images.append(img)
                

        prompt_inputs = self.vlm_module.prepare_model_inputs(
            self.processing_class,
            prompts_text,
            images,
            return_tensors="pt",
            padding=True,
            padding_side="left",
            add_special_tokens=False,
        )
        prompt_inputs = super()._prepare_inputs(prompt_inputs)
        prompt_ids, prompt_mask = prompt_inputs["input_ids"], prompt_inputs["attention_mask"]

[lll6gg]

It seems a single image input. How do you recognise it is multi-images input mode?

[xuexidi]

It seems a single image input. How do you recognise it is multi-images input mode?

1. from VLM-R1 readme：https://github.com/om-ai-lab/VLM-R1?tab=readme-ov-file#multi-image-input

Multi-image Input
We provide an example of multi-image script run_grpo_gui.sh. This task requires the model to analyze two GUI screenshots, taken before and after a user action, to determine if any UI interaction defects are present, which is from GUI-Testing-Arena. Download the image and unzip it into the /path/to/images/. Then modify the image_folders parameter in the script and run it.

bash run_scripts/run_grpo_gui.sh

2. from run_scripts/run_grpo_gui.sh we can see：

https://github.com/om-ai-lab/VLM-R1/blob/c04f284fa2c0f847c9b05f3fd22e6695fb1ec9fe/run_scripts/run_grpo_gui.sh#L30

3. from VLM-R1/src/open-r1-multimodal/src/open_r1/grpo_jsonl.pywe can see：

https://github.com/om-ai-lab/VLM-R1/blob/c04f284fa2c0f847c9b05f3fd22e6695fb1ec9fe/src/open-r1-multimodal/src/open_r1/grpo_jsonl.py#L1036

4. from VLM-R1/src/open-r1-multimodal/src/open_r1/trainer/grpo_trainer.py we can see：

https://github.com/om-ai-lab/VLM-R1/blob/c04f284fa2c0f847c9b05f3fd22e6695fb1ec9fe/src/open-r1-multimodal/src/open_r1/trainer/grpo_trainer.py#L528

[lll6gg]

Well, _get_key_from_inputs ensures it can read an image path or a list of image_paths corresponding to a prompt. They are paired or aligned. Seems no misunderstanding here.