added multi turn token level reward

Box on target: +1.0

All boxes placed: +10.0

Include your thought in <think> </think> tags and your final answer in <answer> </answer> tags.

Your response should be like: <think> [Your thought] </think> <answer> [Your answer] </answer>

Please think step by step and provide the action you want to take.

Your reponse should be in the format of <think>...</think><answer>...</answer>

E.g. <think> There's a box on the upper left of me, I need to push it upward. </think><answer> Right,Up,Up </answer>

"""

# Let's try to use a format reward and answer reward

# If the reponse provides a final answer and is corect, the model receives an accurtacy reward of +1

# is the response encloses its thinking in <think></think> and the final answer is <answer></answer> tags, the model receives a format reward of +1

[Initial Observation]:

{observation}

Decide your next action.

Your reponse should be in the format of <think>...</think><answer>...</answer>

"""

action_template = """After you answer {answer}, the extracted valid action is {valid_action}.\

{observation}

reward: {reward}

done: {done}

Decide your next action.

Your reponse should be in the format of <think>...</think><answer>...</answer>

"""

    rollout_manger.max_turns=1 \

    rollout_manger.window_size=5 \

    trainer.val_before_train=True \

    trainer.val_generations_to_log_to_wandb=5 \

    trainer.val_generations_to_log_to_wandb=2 \

    rollout_manger.n_trajectory=8 \

    2>&1 | tee debug_qwen0_5_1_gpu_grpo.log

    rollout_manger.max_turns=1 \

    rollout_manger.window_size=5 \

    trainer.val_before_train=True \

    trainer.val_generations_to_log_to_wandb=5 \

    trainer.val_generations_to_log_to_wandb=2 \

    rollout_manger.n_trajectory=8 \

    2>&1 | tee debug_qwen0_5_4_gpu_grpo.log

    rollout_manger.max_turns=2 \

    rollout_manger.window_size=5 \

    trainer.val_before_train=True \

    trainer.val_generations_to_log_to_wandb=5 \

    trainer.val_generations_to_log_to_wandb=2 \

    rollout_manger.n_trajectory=8 \

    2>&1 | tee debug_qwen2_5_vl_4gpu_grpo.log

 No newline at end of file

        Returns:

            input_ids: (batch_size, seq_len)

            attention_mask: (batch_size, seq_len)

            loss_mask: (batch_size, seq_len)

            loss_mask: (batch_size, seq_len) # e.g. 0000|1111|0000|11111|000|1111

            token_level_reward_mask: (batch_size, seq_len) # e.g. 0000|0001|0000|00001|000|0001 given the loss mask, mark the position which will be given step reward

        - There will be different stratgy to handel special tokens in the input_ids

        - 1. remove them, in this case we need to fill the hole by adding pad in the right and shift the sequence left

        new_attention_mask = attention_mask.clone()

        loss_mask = torch.zeros_like(input_ids)

        new_loss_mask = torch.zeros_like(input_ids)

        token_level_reward_mask = torch.zeros_like(input_ids)

        new_token_level_reward_mask = torch.zeros_like(input_ids)

        # Process each example in the batch

        for b in range(batch_size):

            # Count right padding tokens using attention mask

            hole_pos=[] # initialize holes position list with last padding token position

            for start_pos, end_pos in zip(sptk_b_indices, sptk_e_indices):

                loss_mask[b][start_pos+1:end_pos] = 1

                token_level_reward_mask[b][end_pos-1] = 1

                hole_pos.append(start_pos.item())

                hole_pos.append(end_pos.item())

            hole_pos.append(seq_len-right_pad_tokens)

                start_pos = hole_pos[i]

                end_pos = hole_pos[i+1]

                new_loss_mask[b,start_pos+1-holes_to_fill:end_pos-holes_to_fill]=loss_mask[b,start_pos+1:end_pos]

                new_token_level_reward_mask[b,start_pos+1-holes_to_fill:end_pos-holes_to_fill]=token_level_reward_mask[b,start_pos+1:end_pos]

                new_input_ids[b,start_pos+1-holes_to_fill:end_pos-holes_to_fill]=input_ids[b,start_pos+1:end_pos]

                new_attention_mask[b,start_pos+1-holes_to_fill:end_pos-holes_to_fill]=attention_mask[b,start_pos+1:end_pos]

                holes_to_fill+=1

            new_input_ids[b][valid_tokens:]=pad_token_id

            new_attention_mask[b][valid_tokens:]=0

        return new_input_ids, new_attention_mask, new_loss_mask

        return new_input_ids, new_attention_mask, new_loss_mask, new_token_level_reward_mask

    def reset(self, env_configs: List[EnvConfig]):

            is_final: Whether the prompt is for the final step 

                - if True, the end of the chat is from the last assistant's response

            prep_for_loss_mask: whether to use special token to wrap llm response

        Returns:

            dict: prompt_with_chat_template : str, image_data: list of images, reward: list of reward

        """

        assert step >= 0

        end_step = step

        assert len(recording) >= end_step + 1, 'History length is not enough'

        history = recording[start_step: end_step + 1]

        rewards=[]

        chat = []

        env_id = history[0]['env_id']

                llm_raw_response = record['info']['llm_raw_response']

                filtered_llm_raw_response = self._handle_special_tokens(llm_raw_response, prep_for_loss_mask=prep_for_loss_mask)

                chat.append({"role": "assistant", "content": filtered_llm_raw_response})

                rewards.append(record['reward'])

            if i<len(history)-1 or not is_final:

                chat.append({"role": "user", "content": record['text_template']})

                if 'image_data' in record:

        return {

            "prompt": prompt_with_chat_template,

            "image_data": image_data,

            "rewards": rewards,

        }

    def _generate_input_item(

        response_rst=self._single_recording_to_prompt(recording, step, window_size, is_final=True, prep_for_loss_mask=True)

        response_with_chat_template=response_rst['prompt']

        image_data=response_rst['image_data']

        rewards=response_rst['rewards']

        has_images = len(image_data) > 0

        row_dict = {}

        attention_mask_prompt=torch.zeros_like(input_ids_prompt) # All prompt will be masked

        input_ids_response, attention_mask_response, loss_mask_response = self._compute_loss_mask(input_ids_response, attention_mask_response)

        input_ids_response, attention_mask_response, loss_mask_response,token_level_reward_mask_response = self._compute_loss_mask(input_ids_response, attention_mask_response)

        input_ids_prompt=input_ids_prompt[0]

        attention_mask_prompt=attention_mask_prompt[0]

        loss_mask_prompt = torch.zeros_like(attention_mask_prompt)

        token_level_reward_mask_prompt = torch.zeros_like(attention_mask_prompt)

        input_ids_response=input_ids_response[0]

        attention_mask_response=attention_mask_response[0]

        loss_mask_response=loss_mask_response[0]

        loss_mask_prompt = torch.zeros_like(attention_mask_prompt)

        token_level_reward_mask_response=token_level_reward_mask_response[0]

        loss_mask = torch.cat([loss_mask_prompt, loss_mask_response], dim=-1)

        token_level_reward_mask = torch.cat([token_level_reward_mask_prompt, token_level_reward_mask_response], dim=-1)

        input_ids = torch.cat([input_ids_prompt, input_ids_response], dim=-1)

        attention_mask = torch.cat([attention_mask_prompt, attention_mask_response], dim=-1)

            delta_position_id = torch.arange(1, response_length + 1, device=position_ids_prompt.device)

            position_ids_response = position_ids_prompt[-1:] + delta_position_id

        reward_positions = torch.nonzero(token_level_reward_mask).squeeze(-1)

        multi_turn_token_level_reward = torch.zeros_like(token_level_reward_mask, dtype=torch.float)

        assert len(reward_positions) == len(rewards), "Number of rewards does not match number of reward positions"

        for idx,reward in enumerate(rewards):

            multi_turn_token_level_reward[reward_positions[idx]] = reward

        position_ids = torch.cat([position_ids_prompt, position_ids_response], dim=-1)

        row_dict['prompts'] = input_ids_prompt

        row_dict['responses'] = input_ids_response

        row_dict['loss_mask'] = loss_mask

        index = row_dict.get("extra_info", {}).get("index", 0)

        row_dict["index"] = index

        row_dict["multi_turn_token_level_reward"] = multi_turn_token_level_reward # (seq_len) later need to convert to (response_len)

        return row_dict

            responses_str = self.tokenizer.batch_decode(

                output_batch.batch['responses'], 

                skip_special_tokens=True

            )

            ) # seems here will remove special token like "<|im_end|>"

            for batch_idx, env_id in self.batch_idx_to_env_id.items(): # TODO whether multiple actions in one rollout are considered here

            for batch_idx, env_id in self.batch_idx_to_env_id.items(): 

                obs, reward, done, info = self.envs[env_id].step(responses_str[batch_idx])

                self.env_states[env_id]['step'] += 1

                self.env_states[env_id]['done'] = done

        inputs=[]

        outputs=[]

        scores=[]

        images=[]

        for k,v in self.recorder.items():

            step=self.env_states[k]['step']

            input_str=self.envs[k].name_repr()+self.envs[k].config_repr(self.envs[k].env_config)

            ouput_rst=self._single_recording_to_prompt(v, step, window_size=None, is_final=True)

            output_str=ouput_rst['prompt']

            ouput_rst=self._single_recording_to_prompt(v, step, window_size=None, is_final=False) # we should see the last obs 

            score=self.envs[k].get_traj_reward()

            inputs.append(input_str)

            outputs.append(output_str)

            outputs.append(ouput_rst['prompt'])

            scores.append(score)

        return inputs,outputs,scores

            images.append(ouput_rst['image_data'])

        return inputs,outputs,scores,images