updated prompts

def free_think_format_prompt(max_actions_per_step, action_sep, add_example=True):

    base_prompt = f"""You can take up to {max_actions_per_step} action(s) at a time, separated by {action_sep}.

You answer should be in the format of:

You should first give your thought process, and then your answer. 

Your response should be in the format of:

<think>...</think><answer>...</answer>"""

    if add_example:

        example = f"\ne.g. <think>I can see the target is on my down left, I should go down then left</think><answer>Down{action_sep}Left</answer>"

        return base_prompt + example

        example = f"""e.g. <think>I can see the target is on my down left, I should go down then left to reach the target</think><answer>Down{action_sep}Left</answer>"""

        return base_prompt + '\n' + example

    return base_prompt

def no_think_format_prompt(max_actions_per_step, action_sep, add_example=True):

    base_prompt = f"""You can take up to {max_actions_per_step} action(s) at a time, separated by {action_sep}.

You answer should be in the format of:

You should provide only your answer.

Your response should be in the format of:

<answer>...</answer>"""

    if add_example:

        example = f"\ne.g. <answer>Down{action_sep}Left</answer>"

        return base_prompt + example

        example = f"""e.g. <answer>Down{action_sep}Left</answer>"""

        return base_prompt + '\n' + example

    return base_prompt

def grounding_format_prompt(max_actions_per_step, action_sep, add_example=True):

    base_prompt = f"""You can take up to {max_actions_per_step} action(s) at a time, separated by {action_sep}.

You answer should be in the format of:

You should first give the current state, then your thought process, and finally your answer.

The state should be in the format of {{"player":(row,column),"target":(row,column)}}

Your response should be in the format of:

<current_state>...</current_state><think>...</think><answer>...</answer>"""

    if add_example:

        example = f"\ne.g. <current_state>I'm in the row 2 col 3. The target is in the row 3 col 2.</current_state><think>I should go down then left to reach the target</think><answer>Down{action_sep}Left</answer>"

        return base_prompt + example

        example = f"""e.g. <current_state>{{"player":(2,3),"target":(3,2)}}</current_state><think>I should go down then left to reach the target</think><answer>Down{action_sep}Left</answer>"""

        return base_prompt + '\n' + example

    return base_prompt

def worldmodeling_format_prompt(max_actions_per_step, action_sep, add_example=True):

    base_prompt = f"""You can take up to {max_actions_per_step} action(s) at a time, separated by {action_sep}.

You answer should be in the format of:

You should first give your thought process, then your answer, and finally predict the next state.

The state should be in the format of {{"player":(row,column),"target":(row,column)}}

Your response should be in the format of:

<think>...</think><answer>...</answer><next_state>...</next_state>"""

    if add_example:

        example = f"\ne.g. <think>I can see the target is on my down left, I should go down then left</think><answer>Down{action_sep}Left</answer><next_state>I'm in the row 3 col 2. The target is in the row 3 col 2.</next_state>"

        return base_prompt + example

        example = f"""e.g. <think>I can see the target is on my down left, I should go down then left</think><answer>Down{action_sep}Left</answer><next_state>{{"player":(3,2),"target":(3,2)}}</next_state>"""

        return base_prompt + '\n' + example

    return base_prompt

def grounding_worldmodeling_format_prompt(max_actions_per_step, action_sep, add_example=True):

    base_prompt = f"""You can take up to {max_actions_per_step} action(s) at a time, separated by {action_sep}.

You answer should be in the format of:

You should first give the current state, then your thought process, then your answer, and finally predict the next state.

The state should be in the format of {{"player":(row,column),"target":(row,column)}}

Your response should be in the format of:

<current_state>...</current_state><think>...</think><answer>...</answer><next_state>...</next_state>"""

    if add_example:

        example = f"\ne.g. <current_state>I'm in the row 2 col 3. The target is in the row 3 col 2.</current_state><think>I should go down then left to reach the target</think><answer>Down{action_sep}Left</answer><next_state>I'm in the row 3 col 2. The target is in the row 3 col 2.</next_state>"

        return base_prompt + example

        example = f"""e.g. <current_state>{{"player":(2,3),"target":(3,2)}}</current_state><think>I should go down then left to reach the target</think><answer>Down{action_sep}Left</answer><next_state>{{"player":(3,2),"target":(3,2)}}</next_state>"""

        return base_prompt + '\n' + example

    return base_prompt

def grounding_symbol_format_prompt(max_actions_per_step, action_sep, add_example=True):

    base_prompt = f"""You can take up to {max_actions_per_step} action(s) at a time, separated by {action_sep}.

You answer should be in the format of:

You should first give the current state as a grid, then your thought process, and finally your answer.

The state should be represented as a grid using the symbols: _ Frozen | O Hole | G Goal | P Player | X Player fell into hole | √ Player on goal.

Your response should be in the format of:

<current_state>...</current_state><think>...</think><answer>...</answer>"""

    if add_example:

        example = f"\ne.g. <current_state>_P__\nG___\n_OO_\n____</current_state><think>I should go down then left to reach the target</think><answer>Down{action_sep}Left</answer>"

        return base_prompt + example

        example = f"""e.g. <current_state>_P__

G___

_OO_

____</current_state><think>I should go down then left to reach the target</think><answer>Down{action_sep}Left</answer>"""

        return base_prompt + '\n' + example

    return base_prompt

def worldmodeling_symbol_format_prompt(max_actions_per_step, action_sep, add_example=True):

    base_prompt = f"""You can take up to {max_actions_per_step} action(s) at a time, separated by {action_sep}.

You answer should be in the format of:

You should first give your thought process, then your answer, and finally predict the next state as a grid.

The state should be represented as a grid using the symbols: _ Frozen | O Hole | G Goal | P Player | X Player fell into hole | √ Player on goal.

Your response should be in the format of:

<think>...</think><answer>...</answer><next_state>...</next_state>"""

    if add_example:

        example = f"\ne.g. <think>I can see the target is on my down left, I should go down then left</think><answer>Down{action_sep}Left</answer><next_state>____\n√___\n_OO_\n____</next_state>"

        return base_prompt + example

        example = f"""e.g. <think>I can see the target is on my down left, I should go down then left</think><answer>Down{action_sep}Left</answer><next_state>____

√___

_OO_

____</next_state>"""

        return base_prompt + '\n' + example

    return base_prompt

def grounding_worldmodeling_symbol_format_prompt(max_actions_per_step, action_sep, add_example=True):

    base_prompt = f"""You can take up to {max_actions_per_step} action(s) at a time, separated by {action_sep}.

You answer should be in the format of:

You should first give the current state as a grid, then your thought process, then your answer, and finally predict the next state as a grid.

The state should be represented as grids using the symbols: _ Frozen | O Hole | G Goal | P Player | X Player fell into hole | √ Player on goal.

Your response should be in the format of:

<current_state>...</current_state><think>...</think><answer>...</answer><next_state>...</next_state>"""

    if add_example:

        example = f"\ne.g. <current_state>_P__\nG___\n_OO_\n____</current_state><think>I should go down then left to reach the target</think><answer>Down{action_sep}Left</answer><next_state>____\n√___\n_OO_\n____</next_state>"

        return base_prompt + example

        example = f"""e.g. <current_state>_P__

G___

_OO_

____</current_state><think>I should go down then left to reach the target</think><answer>Down{action_sep}Left</answer><next_state>____

√___

_OO_

____</next_state>"""

        return base_prompt + '\n' + example

    return base_prompt

# Dictionary mapping format names to their corresponding functions

def free_think_format_prompt(max_actions_per_step, action_sep, add_example=True):

    base_prompt = f"""You can take up to {max_actions_per_step} action(s) at a time, separated by {action_sep}.

Your answer should be in the format of:

You should first give your thought process, and then your answer.

Your response should be in the format of:

<think>...</think><answer>...</answer>"""

    if add_example:

        example = f"\ne.g. <think>The box is one step below me, and the target is two steps below me, I need to go down then push the box down to the target.</think><answer>Down{action_sep}Down</answer>"

        return base_prompt + example

        example = f"""e.g. <think>The box is one step below me, and the target is two steps below me, I need to go down then push the box down to the target.</think><answer>Down{action_sep}Down</answer>"""

        return base_prompt + '\n' + example

    return base_prompt

def no_think_format_prompt(max_actions_per_step, action_sep, add_example=True):

    base_prompt = f"""You can take up to {max_actions_per_step} action(s) at a time, separated by {action_sep}.

Your answer should be in the format of:

You should provide only your answer.

Your response should be in the format of:

<answer>...</answer>"""

    if add_example:

        example = f"\ne.g. <answer>Down{action_sep}Down</answer>"

        return base_prompt + example

        example = f"""e.g. <answer>Down{action_sep}Down</answer>"""

        return base_prompt + '\n' + example

    return base_prompt

def grounding_format_prompt(max_actions_per_step, action_sep, add_example=True):

    base_prompt = f"""You can take up to {max_actions_per_step} action(s) at a time, separated by {action_sep}.

Your answer should be in the format of:

You should first give the current state, then your thought process, and finally your answer.

The state should be in the format of {{"player":(row,column),"box":(row1,column1),"target":(row1,column1)}}

Your response should be in the format of:

<current_state>...</current_state><think>...</think><answer>...</answer>"""

    if add_example:

        example = f"\ne.g. <current_state>I'm at row 2 col 3. The box is at row 4 col 3. The target is at at row 5 col 3.</current_state><think>I need to go down then push the box down to the target</think><answer>Down{action_sep}Down</answer>"

        return base_prompt + example

        example = f"""e.g. <current_state>{{"player":(2,3),"box":(4,3),"target":(5,3)}}</current_state><think>I need to go down then push the box down to the target</think><answer>Down{action_sep}Down</answer>"""

        return base_prompt + '\n' + example

    return base_prompt

def worldmodeling_format_prompt(max_actions_per_step, action_sep, add_example=True):

    base_prompt = f"""You can take up to {max_actions_per_step} action(s) at a time, separated by {action_sep}.

Your answer should be in the format of:

You should first give your thought process, then your answer, and finally predict the next state.

The state should be in the format of {{"player":(row,column),"box":(row1,column1),"target":(row1,column1)}}

Your response should be in the format of:

<think>...</think><answer>...</answer><next_state>...</next_state>"""

    if add_example:

        example = f"\ne.g. <think>I need to go down then push the box down to the target.</think><answer>Down{action_sep}Down</answer><next_state>I'm at row 4 col 3. The box is at row 5 col 3. The target is at at row 5 col 3.</next_state>"

        return base_prompt + example

        example = f"""e.g. <think>I need to go down then push the box down to the target.</think><answer>Down{action_sep}Down</answer><next_state>{{"player":(4,3),"box":(5,3),"target":(5,3)}}</next_state>"""

        return base_prompt + '\n' + example

    return base_prompt

def grounding_worldmodeling_format_prompt(max_actions_per_step, action_sep, add_example=True):

    base_prompt = f"""You can take up to {max_actions_per_step} action(s) at a time, separated by {action_sep}.

Your answer should be in the format of:

You should first give the current state, then your thought process, then your answer, and finally predict the next state.

The state should be in the format of {{"player":(row,column),"box":(row1,column1),"target":(row1,column1)}}

Your response should be in the format of:

<current_state>...</current_state><think>...</think><answer>...</answer><next_state>...</next_state>"""

    if add_example:

        example = f"\ne.g. <current_state>I'm at row 2 col 3. The box is at row 4 col 3. The target is at at row 5 col 3.</current_state><think>I need to go down then push the box down to the target</think><answer>Down{action_sep}Down</answer><next_state>I'm at row 4 col 3. The box is at row 5 col 3. The target is at at row 5 col 3.</next_state>"

        return base_prompt + example

        example = f"""e.g. <current_state>{{"player":(2,3),"box":(4,3),"target":(5,3)}}</current_state><think>I need to go down then push the box down to the target</think><answer>Down{action_sep}Down</answer><next_state>{{"player":(4,3),"box":(5,3),"target":(5,3)}}</next_state>"""

        return base_prompt + '\n' + example

    return base_prompt

def grounding_symbol_format_prompt(max_actions_per_step, action_sep, add_example=True):

    base_prompt = f"""Symbols:

# Wall | _ Floor | O Target | X Box | P You | √ Box on Target | S You on Target

You can take up to {max_actions_per_step} action(s) at a time, separated by {action_sep}.

Your answer should be in the format of:

    base_prompt = f"""You can take up to {max_actions_per_step} action(s) at a time, separated by {action_sep}.

You should first give the current state as a grid, then your thought process, and finally your answer.

The state should be represented as a grid using the symbols: # Wall | _ Floor | O Target | X Box | P You | √ Box on Target | S You on Target.

Your response should be in the format of:

<current_state>...</current_state><think>...</think><answer>...</answer>"""

    if add_example:

        example = f"\ne.g. <current_state>####\n#_P#\n#__#\n#_x#\n#_O#</current_state><think>I need to go down then push the box down to reach the target</think><answer>Down{action_sep}Down</answer>"

        return base_prompt + example

        example = f"""e.g. <current_state>####

#_P#

#__#

#_X#

#_O#</current_state><think>I need to go down then push the box down to reach the target</think><answer>Down{action_sep}Down</answer>"""

        return base_prompt + '\n' + example

    return base_prompt

def worldmodeling_symbol_format_prompt(max_actions_per_step, action_sep, add_example=True):

    base_prompt = f"""You can take up to {max_actions_per_step} action(s) at a time, separated by {action_sep}.

Your answer should be in the format of:

You should first give your thought process, then your answer, and finally predict the next state as a grid.

The state should be represented as a grid using the symbols: # Wall | _ Floor | O Target | X Box | P You | √ Box on Target | S You on Target.

Your response should be in the format of:

<think>...</think><answer>...</answer><next_state>...</next_state>"""

    if add_example:

        example = f"\ne.g. <think>I need to go down then push the box down to reach the target</think><answer>Down{action_sep}Down</answer><next_state>####\n#__#\n#__#\n#_P#\n#_√#</next_state>"

        return base_prompt + example

        example = f"""e.g. <think>I need to go down then push the box down to reach the target</think><answer>Down{action_sep}Down</answer><next_state>####

#__#

#__#

#_P#

#_√#</next_state>"""

        return base_prompt + '\n' + example

    return base_prompt

def grounding_worldmodeling_symbol_format_prompt(max_actions_per_step, action_sep, add_example=True):

    base_prompt = f"""You can take up to {max_actions_per_step} action(s) at a time, separated by {action_sep}.

Your answer should be in the format of:

You should first give the current state as a grid, then your thought process, then your answer, and finally predict the next state as a grid.

The state should be represented as grids using the symbols: # Wall | _ Floor | O Target | X Box | P You | √ Box on Target | S You on Target.

Your response should be in the format of:

<current_state>...</current_state><think>...</think><answer>...</answer><next_state>...</next_state>"""

    if add_example:

        example = f"\ne.g. <current_state>####\n#_P#\n#__#\n#_x#\n#_O#</current_state><think>I need to go down then push the box down to reach the target</think><answer>Down{action_sep}Down</answer><next_state>####\n#__#\n#__#\n#_P#\n#_√#</next_state>"

        return base_prompt + example

        example = f"""e.g. <current_state>####

#_P#

#__#

#_X#

#_O#</current_state><think>I need to go down then push the box down to reach the target</think><answer>Down{action_sep}Down</answer><next_state>####

#__#

#__#

#_P#

#_√#</next_state>"""

        return base_prompt + '\n' + example

    return base_prompt

# Dictionary mapping format names to their corresponding functions