updated prompt formats (1666f19f) · Commits · 张泽凯 / VAGEN

vagen/env/navigation/env.py

+56 −26

Original line number	Diff line number	Diff line
		@@ -3,18 +3,17 @@ import ai2thor.controller
		import numpy as np
		import time
		import math
		import re
		from ai2thor.platform import CloudRendering
		from typing import Dict, List, Tuple, Optional, Any
		from vagen.env.utils.env_utils import NoLoggerWarnings, set_seed
		from vagen.env.utils.context_utils import parse_llm_raw_response, convert_numpy_to_PIL
		from vagen.env.utils.parse_utils import parse_function_map
		from .env_config import NavigationEnvConfig
		from .prompt import system_prompt_text, system_prompt_vision, init_observation_template, action_template
		from .prompt import system_prompt, system_prompt_vision, init_observation_template, action_template, format_prompt


		class NavigationEnv(BaseEnv):
		"""Navigation environment based on AI2-THOR."""

		"""Navigation environment from embodied bench. """
		SUCCESS_THRESHOLD = 1

		ValidEvalSets = [
		@@ -49,7 +48,8 @@ class NavigationEnv(BaseEnv):
		"""Initialize the Navigation environment.

		Args:
		config: Configuration for the environment
		config: Configuration for the environment including resolution, FOV,
		eval set, render mode, etc.
		"""
		super().__init__()
		self.config = config
		@@ -70,12 +70,7 @@ class NavigationEnv(BaseEnv):
		}

		# Initialize AI2-THOR controller
		# try:
		self.env = ai2thor.controller.Controller(**self.thor_config)
		# except Exception as e:
		# print(f"Error initializing AI2-THOR: {e}")
		# import traceback
		# traceback.print_exc()

		# Load dataset
		assert config.eval_set in self.ValidEvalSets
		@@ -98,6 +93,14 @@ class NavigationEnv(BaseEnv):
		self.multiview = config.multiview
		self.img_paths = []
		self.total_reward = 0
		self.valid_actions = []
		self.reward = 0

		# Store the format prompt function for later use
		self.format_prompt_func = format_prompt[self.config.prompt_format]

		# Get the parse function based on the prompt format
		self.parse_func = parse_function_map[self.config.prompt_format]

		def _get_dataset_path(self, eval_set):
		"""Get the path to the dataset file."""
		@@ -118,13 +121,18 @@ class NavigationEnv(BaseEnv):
		def reset(self, seed=None):
		"""Reset the environment to a new episode.

		This method resets the AI2-THOR environment and initializes a new episode
		based on the dataset. If a seed is provided, it ensures deterministic
		episode selection.

		Args:
		seed: Random seed for reproducibility

		Returns:
		Observation dict, info dict
		"""
		# Reset the environment
		# Reset the environment with the proper seed

		idx = seed % self.number_of_episodes if seed is not None else 0

		# Get the trajectory data
		@@ -173,13 +181,20 @@ class NavigationEnv(BaseEnv):
		self._episode_start_time = time.time()
		self.img_paths = []
		self.total_reward = 0

		self.valid_actions = []
		self.reward = 0

		return self._render(init_obs=True), {}

		def step(self, action_str: str):
		"""Execute an action in the environment.

		This method:
		1. Parses the raw LLM response to extract actions
		2. Executes each valid action in sequence
		3. Calculates rewards and metrics
		4. Generates the next observation

		Args:
		action_str: Raw text response from LLM

		@@ -187,7 +202,7 @@ class NavigationEnv(BaseEnv):
		Observation, reward, done, info
		"""
		# Process the LLM response to extract actions
		rst = parse_llm_raw_response(
		rst = self.parse_func(
		response=action_str,
		special_token_list=self.config.get('special_token_list', None),
		action_sep=self.config.get('action_sep', ','),
		@@ -215,6 +230,8 @@ class NavigationEnv(BaseEnv):

		# Execute valid actions
		if metrics["turn_metrics"]["action_is_valid"]:
		# Add format reward if actions were valid and format is correct
		if rst.get("format_correct", True):
		self.reward += self.config.format_reward

		for action in action_list:
		@@ -226,7 +243,6 @@ class NavigationEnv(BaseEnv):

		# Update reward based on success
		if success:
		# Success reward is handled in the measure_success method
		self.reward += 10.0 # Success reward
		done = True
		metrics['traj_metrics']['success'] = True
		@@ -305,6 +321,10 @@ class NavigationEnv(BaseEnv):
		def _render(self, init_obs=False):
		"""Render the environment observation.

		This method creates either a text representation or an image of the environment
		state, depending on the configured render mode. It formats the observation string
		based on whether this is the initial observation or a subsequent one.

		Args:
		init_obs: Whether this is the initial observation

		@@ -313,6 +333,13 @@ class NavigationEnv(BaseEnv):
		"""
		img_placeholder = self.config.get("image_placeholder", "<image>")

		# Get format prompt without examples for action/init templates
		format_prompt_text = self.format_prompt_func(
		max_actions_per_step=self.config.max_actions_per_step,
		action_sep=self.config.action_sep,
		add_example=False # No examples for action and init obs
		)

		# Get the RGB frame from the environment
		frame = self.env.last_event.frame

		@@ -323,18 +350,18 @@ class NavigationEnv(BaseEnv):

		# Format the template
		if init_obs:
		obs_str = init_observation_template.format(
		obs_str = init_observation_template(
		observation=img_placeholder,
		instruction=self.episode_language_instruction,
		)
		) + "\n" + format_prompt_text
		else:
		obs_str = action_template.format(
		obs_str = action_template(
		valid_action=self.valid_actions,
		observation=img_placeholder,
		reward=self.reward,
		done=self.measure_success()[0],
		instruction=self.episode_language_instruction,
		)
		) + "\n" + format_prompt_text

		return {
		"obs_str": obs_str,
		@@ -344,19 +371,23 @@ class NavigationEnv(BaseEnv):
		def system_prompt(self):
		"""Get the system prompt for the environment.

		Returns a prompt explaining the environment to the LLM agent,
		with different prompts for text and vision modes.

		Returns:
		System prompt string
		"""
		if self.config.render_mode == "vision":
		return system_prompt_vision.format(
		# Get format prompt with examples for system prompt
		format_prompt_text = self.format_prompt_func(
		max_actions_per_step=self.config.max_actions_per_step,
		action_sep=self.config.action_sep
		action_sep=self.config.action_sep,
		add_example=True # Always true for system prompt
		)

		if self.config.render_mode == "vision":
		return system_prompt_vision() + '\n' + format_prompt_text
		else:
		return system_prompt_text.format(
		max_actions_per_step=self.config.max_actions_per_step,
		action_sep=self.config.action_sep
		)
		return system_prompt() + '\n' + format_prompt_text

		def compute_reward(self):
		"""Compute the total reward for the episode.
		@@ -386,7 +417,6 @@ if __name__ == "__main__":
		img.save(f"./test_navigation/navigation_{i}.png")
		done = False


		# Interactive testing loop
		while not done:
		i += 1

vagen/env/navigation/env_config.py

+3 −0

Original line number	Diff line number	Diff line
		@@ -15,6 +15,9 @@ class NavigationEnvConfig(BaseEnvConfig):
		max_action_penalty: float = -0.1
		format_reward: float = 0.5
		gpu_device: int = 0
		prompt_format: str = "free_think"
		# "free_think", "no_think", "grounding", "worldmodeling", "grounding_worldmodeling"
		use_accuracy_reward: bool = False

		def config_id(self) -> str:
		"""Generate a unique identifier for this configuration."""

vagen/env/navigation/prompt.py

+75 −17

Original line number	Diff line number	Diff line
		system_prompt_text = """You are a home robot and perform navigation tasks according to instructions.
		def system_prompt():
		return """You are a home robot and perform navigation tasks according to instructions.

		Navigation Guide
		Goal: Achieve the human instruction
		@@ -17,22 +18,15 @@ lookdown: Tilt the camera downward by 30 degrees
		Rewards:
		Format correct: +0.5
		Achieve the human instruction: +10.0

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

		You can take up to {max_actions_per_step} action(s) at a time, separated by {action_sep}.
		Your reponse should be in the format of <think>...</think><answer>...</answer>
		e.g.
		<think>I can see from the sight the target object is right in the top left of me, I will move forward, then move left to access it.</think><answer>moveahead{action_sep}moveahead{action_sep}moveahead{action_sep}moveahead{action_sep}moveahead{action_sep}moveleft{action_sep}moveleft</answer>
		"""

		system_prompt_vision = """You are a home robot and perform navigation tasks according to instructions.
		def system_prompt_vision():
		return """You are a home robot and perform navigation tasks according to instructions.

		Navigation Guide
		You should follow the human instruction and navigate to the target location.

		Actions you can take: moveahead, moveback, moveright, moveleft, rotateright, rotateleft, lookup, lookdown.
		You can take up to {max_actions_per_step} action(s) at a time.

		moveahead: Move forward by 0.25 meter
		moveback: Move backward by 0.25 meter
		@@ -48,24 +42,88 @@ Format correct: +0.5
		Achieve the human instruction: +10.0

		The instruction will be provided with each observation. Look at the image carefully and navigate to complete the instruction.
		e.g.
		<think>I can see from the sight the target object is right in the top left of me, I will move forward, then move left to access it.</think><answer>moveahead, moveahead,moveahead,moveahead,moveahead,moveleft,moveleft</answer>
		"""

		init_observation_template = """
		[Initial Observation]:
		def init_observation_template(observation, instruction):
		return f"""[Initial Observation]:
		{observation}
		Human Instruction: {instruction}
		Decide your next action(s).
		Your reponse should be in the format of <think>...</think><answer>...</answer>
		"""

		action_template = """After your answer, the extracted valid action is {valid_action}.
		def action_template(valid_action, observation, reward, done, instruction):
		return f"""After your answer, the extracted valid action is {valid_action}.
		After that, the observation is:
		{observation}
		reward: {reward}
		done: {done}
		Human Instruction: {instruction}
		Decide your next action(s).
		Your reponse should be in the format of <think>...</think><answer>...</answer>
		"""

		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 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"""e.g. <think>I can see from the sight the target object is right in the top left of me, I will move forward, then move left to access it.</think><answer>moveahead{action_sep}moveahead{action_sep}moveahead{action_sep}moveahead{action_sep}moveahead{action_sep}moveleft{action_sep}moveleft</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 should provide only your answer.
		Your response should be in the format of:
		<answer>...</answer>"""

		if add_example:
		example = f"""e.g. <answer>moveahead{action_sep}moveahead{action_sep}moveahead{action_sep}moveahead{action_sep}moveahead{action_sep}moveleft{action_sep}moveleft</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 should first give the current state, then your thought process, and finally your answer.
		The state should be described in detail about what you see in the environment.
		Your response should be in the format of:
		<current_state>...</current_state><think>...</think><answer>...</answer>"""

		if add_example:
		example = f"""e.g. <current_state>I am in a living room. There is a couch to my left, a TV in front of me, and a doorway to the kitchen on my right. The target object, a vase, appears to be on a shelf near the kitchen doorway.</current_state><think>I need to move toward the kitchen doorway to reach the vase. I'll move forward to get closer to the center of the room, then turn right and move toward the kitchen.</think><answer>moveahead{action_sep}moveahead{action_sep}rotateright{action_sep}moveahead{action_sep}moveahead</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 should first give your thought process, then your answer, and finally predict the next state.
		The next state should describe what you expect to see after your actions are executed.
		Your response should be in the format of:
		<think>...</think><answer>...</answer><next_state>...</next_state>"""

		if add_example:
		example = f"""e.g. <think>I can see the kitchen doorway to my right, and I need to go there to find the refrigerator. I'll turn right and move forward.</think><answer>rotateright{action_sep}moveahead{action_sep}moveahead</answer><next_state>I am now in the kitchen doorway. In front of me is the kitchen counter with a sink. To the left I can see a refrigerator against the wall. There's a kitchen island in the center of the room.</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 should first give the current state, then your thought process, then your answer, and finally predict the next state.
		Both the current and next states should describe what you see or expect to see in the environment.
		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"""e.g. <current_state>I am at the entrance of a bedroom. There is a bed to the left, a desk with a lamp on the right, and a closet straight ahead. The target object, a book, appears to be on the desk.</current_state><think>I need to move toward the desk to reach the book. I'll turn right and move forward.</think><answer>rotateright{action_sep}moveahead{action_sep}moveahead</answer><next_state>I am now standing in front of the desk. The desk has a lamp, a computer, and several books on it. The target book is within reach on the right side of the desk.</next_state>"""
		return base_prompt + '\n' + example
		return base_prompt

		# Dictionary mapping format names to their corresponding functions
		format_prompt = {
		"free_think": free_think_format_prompt,
		"no_think": no_think_format_prompt,
		"grounding": grounding_format_prompt,
		"worldmodeling": worldmodeling_format_prompt,
		"grounding_worldmodeling": grounding_worldmodeling_format_prompt
		}
		No newline at end of file

Admin message