updated maniskill for different prompting format (b4bf55aa) · Commits · 张泽凯 / VAGEN

vagen/env/primitive_skill/env.py

+136 −19

Original line number	Diff line number	Diff line
		@@ -3,14 +3,22 @@ import numpy as np
		import copy
		from typing import Dict, List, Optional, Tuple, Any
		from gymnasium.utils import seeding
		from vagen.env.utils.context_utils import parse_llm_raw_response, convert_numpy_to_PIL
		from vagen.env.utils.context_utils import convert_numpy_to_PIL
		from vagen.env.utils.parse_utils import parse_function_map
		from .env_config import PrimitiveSkillEnvConfig
		from .maniskill.utils import build_env, handle_info, get_workspace_limits
		from .prompts import system_prompt, init_observation_template, action_template
		from .prompts import system_prompt, init_observation_template, action_template, format_prompt
		import vagen.env.primitive_skill.maniskill.env

		class PrimitiveSkillEnv(BaseEnv):
		def __init__(self, config: PrimitiveSkillEnvConfig):
		"""
		Initialize the PrimitiveSkill environment.

		Args:
		config (PrimitiveSkillEnvConfig): Configuration parameters for the environment
		"""
		BaseEnv.__init__(self)
		self.config = config
		if self.config.record_video:
		record_dir = self.config.video_record_dir
		@@ -18,7 +26,25 @@ class PrimitiveSkillEnv(BaseEnv):
		record_dir = None
		self.env = build_env(config.env_id, record_dir=record_dir)

		# Store the format prompt function for later use based on the configuration
		self.format_prompt_func = format_prompt[self.config.prompt_format]
		self.parse_func = parse_function_map[self.config.prompt_format]
		# Define the state keys for the environment
		self.state_keys = self.env.state_keys

		def reset(self, seed: Optional[int] = None) -> Tuple[Dict[str, Any], Dict[str, Any]]:
		"""
		Reset the environment to an initial state.

		Args:
		seed (Optional[int]): Random seed for environment generation
		If None, a random seed is used

		Returns:
		Tuple[Dict, Dict]:
		- obs: Dictionary containing observation string and optional image data
		- info: Empty dictionary for initial state
		"""
		_, info = self.env.reset(seed=seed)
		obs = self._render(info, init_obs=True)
		self.last_info = info
		@@ -28,13 +54,25 @@ class PrimitiveSkillEnv(BaseEnv):
		return obs, {}

		def step(self, action_str):
		"""
		Take a step in the environment based on the agent's action.

		Args:
		action_str (str): Raw string from LLM containing actions

		Returns:
		Tuple[Dict, float, bool, Dict]:
		- obs: Dictionary with observation string and optional image data
		- reward: Numeric reward for the step
		- done: Boolean indicating if episode is complete
		- info: Dictionary containing metrics and parsed action data
		"""
		reward = 0
		rst = parse_llm_raw_response(
		response=action_str,
		rst = self.parse_func( response=action_str,
		special_token_list=self.config.special_token_list,
		action_sep=self.config.action_sep,
		max_actions=self.config.max_actions_per_step,
		)
		max_actions=self.config.max_actions_per_step)

		output_info = {}
		output_info.update(rst)
		valid_actions = []
		@@ -46,8 +84,12 @@ class PrimitiveSkillEnv(BaseEnv):
		"success": False, # Will be set to True if agent reaches goal
		},
		}

		info = self.last_info
		terminated, truncated = False, False


		# Execute each action in the list
		for action in rst['actions']:
		parsed_action = self._parse_action(action)
		if parsed_action is not None:
		@@ -61,30 +103,57 @@ class PrimitiveSkillEnv(BaseEnv):
		break
		if truncated or terminated:
		break

		# Check if actions were valid and format was correct
		metrics["turn_metrics"]['action_is_valid'] = len(valid_actions) > 0 and len(valid_actions) == len(rst['actions'])
		if metrics["turn_metrics"]['action_is_valid']:
		if metrics["turn_metrics"]['action_is_valid'] and rst["format_correct"]:
		reward += self.config.format_reward
		if info['is_success']:
		# Check for success
		if info.get('is_success', False):
		metrics["traj_metrics"]['success'] = True

		done = terminated or truncated
		info["action_is_valid"] = metrics["turn_metrics"]['action_is_valid']

		obs = self._render(info, init_obs=False, valid_actions=valid_actions)
		output_info["metrics"] = metrics

		self.total_reward += reward
		if isinstance(done, np.ndarray):
		done = done.item()

		return obs, reward, done, output_info

		def system_prompt(self):
		return system_prompt.format(
		"""
		Get the system prompt for the environment.

		Returns:
		str: System prompt string with environment description and instructions
		"""
		# 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,
		state_keys=self.state_keys,
		add_example=True # Always true for system prompt
		)

		return system_prompt() + '\n' + format_prompt_text

		def close(self):
		"""
		Close the environment and clean up resources.
		"""
		self.env.close()

		def _compute_reward(self):
		"""
		Calculate the reward based on environment state.

		Returns:
		float: Computed reward value
		"""
		if self.last_info.get("success", False):
		return 10

		@@ -104,10 +173,36 @@ class PrimitiveSkillEnv(BaseEnv):
		return (max_stage + 1) * 2

		def compute_reward(self):
		"""
		Get the cumulative reward for the episode.

		Returns:
		float: Total reward accumulated during the current episode
		"""
		return self._compute_reward() + self.total_reward - self.initial_reward - self.steps * 0.1

		def _get_current_state(self):
		"""
		Get a representation of the current state for comparison.

		Returns:
		dict: Dictionary representation of important state components
		"""
		# This is a simple implementation - customize based on your environment
		return {k: v for k, v in self.last_info.items() if k.endswith('_position')}

		def _render(self, info, init_obs=False, valid_actions=None):
		"""
		Render the environment as an observation.

		Args:
		info (dict): Environment info dictionary
		init_obs (bool): If True, create initial observation
		valid_actions (list): List of valid actions executed (for step observations)

		Returns:
		Dict: Observation dictionary containing observation string and optional image data
		"""
		new_info = handle_info(info.copy(), mask_success=self.config.mask_success, env=self.env)
		object_positions = new_info['obj_positions']
		other_information = new_info['other_info']
		@@ -115,31 +210,46 @@ class PrimitiveSkillEnv(BaseEnv):
		img_placeholder = self.config.image_placeholder
		x_workspace, y_workspace, z_workspace = get_workspace_limits(self.env)

		# 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,
		state_keys=self.state_keys,
		add_example=False # No examples for action and init obs
		)

		if init_obs:
		obs_str = init_observation_template.format(observation=img_placeholder,
		# Initial observation
		obs_str = init_observation_template(
		observation=img_placeholder,
		instruction=instruction,
		object_positions=object_positions,
		other_information=other_information,
		x_workspace=x_workspace,
		y_workspace=y_workspace,
		z_workspace=z_workspace,
		max_action=self.config.max_actions_per_step)
		object_positions=object_positions,
		other_information=other_information
		) + "\n" + format_prompt_text
		else:
		obs_str = action_template.format(valid_actions=valid_actions,
		# Subsequent observations include action results
		obs_str = action_template(
		valid_actions=valid_actions,
		observation=img_placeholder,
		instruction=instruction,
		object_positions=object_positions,
		other_information=other_information,
		x_workspace=x_workspace,
		y_workspace=y_workspace,
		z_workspace=z_workspace,
		max_action=self.config.max_actions_per_step)
		object_positions=object_positions,
		other_information=other_information
		) + "\n" + format_prompt_text

		multi_modal_data = None
		if self.config.render_mode == "vision":
		img = self.env.render()
		multi_modal_data = {
		img_placeholder: [convert_numpy_to_PIL(img)]
		}

		# Return observation dictionary with appropriate fields
		if multi_modal_data is not None:
		return {
		"obs_str": obs_str,
		@@ -150,9 +260,17 @@ class PrimitiveSkillEnv(BaseEnv):
		"obs_str": obs_str,
		}

		import numpy as np

		def _parse_action(self, action_str):
		"""
		Parse a single action string into an action array.

		Args:
		action_str (str): Action string to parse

		Returns:
		np.array: Parsed action array or None if invalid
		"""
		# Initialize empty 9-dim array (3 for action type, 6 for coordinates)
		action_array = np.zeros(9)

		@@ -210,7 +328,6 @@ class PrimitiveSkillEnv(BaseEnv):
		# If any parsing error occurs, return None
		return None


		if __name__ == "__main__":
		"""
		Example usage of the manipulation environment.

vagen/env/primitive_skill/env_config.py

+4 −0

Original line number	Diff line number	Diff line
		@@ -11,6 +11,10 @@ class PrimitiveSkillEnvConfig(BaseEnvConfig):
		record_video: bool = field(default=False)
		video_record_dir: str = field(default='./test')
		mask_success: bool = field(default=True)
		prompt_format: str = "free_think"
		# "free_think", "no_think", "grounding", "worldmodeling", "grounding_worldmodeling"
		use_accuracy_reward: bool = False


		def config_id(self) -> str:
		id_fields=["env_id","render_mode","max_actions_per_step"]

vagen/env/primitive_skill/prompts.py

+143 −15

Original line number	Diff line number	Diff line
		system_prompt = """You are an AI assistant controlling a Franka Emika robot arm. Your goal is to understand human instructions and translate them into a sequence of executable actions for the robot, based on visual input and the instruction.
		def system_prompt():
		"""
		Returns the system prompt for the robot arm control.

		Returns:
		str: The system prompt
		"""
		return """You are an AI assistant controlling a Franka Emika robot arm. Your goal is to understand human instructions and translate them into a sequence of executable actions for the robot, based on visual input and the instruction.

		Action Space Guide
		You can command the robot using the following actions:
		@@ -10,16 +17,12 @@ You can command the robot using the following actions:
		Hints:
		1. The coordinates (x, y, z) are in millimeters and are all integers.
		2. Please ensure that the coordinates are within the workspace limits.
		3. The position is the center of the object, when you place, please consider the volume of the obeject. It's always fine to set z much higher.

		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 need to pick obj A (100,100,100) first and place it at the obj B (200,200,200)</think><answer>pick(100,100,100)\|place(200,200,400)</answer>
		e.g. <think>I should push obj A (100,200,20) along the y axis</think><answer>push(100,200,20,100,400,20)</answer>
		3. The position is the center of the object, when you place, please consider the volume of the object. It's always fine to set z much higher when placing an item.
		4. We will provide the object positions to you, but you need to match them to the object in the image by yourself. You're facing toward the negative x-axis, and the negative y-axis is to your left, the positive y-axis is to your right, and the positive z-axis is up.
		"""

		init_observation_template = """
		def init_observation_template(observation, instruction, x_workspace, y_workspace, z_workspace, object_positions, other_information):
		return f"""
		[Initial Observation]:
		{observation}
		Human Instruction: {instruction}
		@@ -30,11 +33,11 @@ Object positions:
		{object_positions}
		Other information:
		{other_information}
		Decide your next action(s), you can propose at most {max_action} actions.
		Your reponse should be in the format of <think>...</think><answer>...</answer>
		"""
		Decide your next action(s)."""

		def action_template(valid_actions, observation, instruction, x_workspace, y_workspace, z_workspace, object_positions, other_information):

		action_template = """After your answer, the extracted valid action(s) is {valid_actions}.
		return f"""After your answer, the extracted valid action(s) is {valid_actions}.
		After that, the observation is:
		{observation}
		Human Instruction: {instruction}
		@@ -45,6 +48,131 @@ Object positions:
		{object_positions}
		Other information:
		{other_information}
		Decide your next action(s), you can propose at most {max_action} actions.
		Your reponse should be in the format of <think>...</think><answer>...</answer>
		Decide your next action(s)."""

		def free_think_format_prompt(max_actions_per_step, action_sep, state_keys, add_example=True):
		"""
		Format prompt for free thinking: thinking + answer.

		Args:
		max_actions_per_step (int): Maximum number of actions allowed per step
		action_sep (str): Separator between actions
		state_keys (list): List of object states to track/predict (not used in this format)
		add_example (bool): Whether to add an example

		Returns:
		str: The formatted prompt
		"""
		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 need to pick the red cube at (100,100,40) first and place it on top of the green cube at (200,200,50)</think><answer>pick(100,100,40){action_sep}place(200,200,100)</answer>"""
		return base_prompt + '\n' + example
		return base_prompt

		def no_think_format_prompt(max_actions_per_step, action_sep, state_keys, 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>pick(100,100,40){action_sep}place(200,200,100)</answer>"""
		return base_prompt + '\n' + example
		return base_prompt

		def grounding_format_prompt(max_actions_per_step, action_sep, state_keys, add_example=True):
		state_format = {key: "(x,y,z)" for key in state_keys}

		# Create example state with first object at pick position and others at different positions
		state_example = {}
		for i, key in enumerate(state_keys):
		if i == 0:
		state_example[key] = "(100,100,40)" # This will be the object to pick
		else:
		state_example[key] = f"({i100+200},{i100+200},{50})"

		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 in the format of {state_format}
		Your response should be in the format of:
		<current_state>...</current_state><think>...</think><answer>...</answer>"""

		if add_example:
		# Use the first key as the object being manipulated in the example
		target_object = state_keys[0] if state_keys else "red_cube_position"
		example = f"""e.g. <current_state>{state_example}</current_state><think>I need to pick the {target_object.replace('_position','')} at (100,100,40) and place it at a new location</think><answer>pick(100,100,40){action_sep}place(80,100,50)</answer>"""
		return base_prompt + '\n' + example
		return base_prompt

		def worldmodeling_format_prompt(max_actions_per_step, action_sep, state_keys, add_example=True):
		state_format = {key: "(x,y,z)" for key in state_keys}

		# Set up the next state example showing the object movement
		next_state_example = {}

		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 state should be in the format of {state_format}
		Your response should be in the format of:
		<think>...</think><answer>...</answer><next_state>...</next_state>"""

		if add_example:
		# Use the first key as the object being manipulated in the example
		target_object = state_keys[0] if state_keys else "red_cube_position"

		# Create next state example showing the object moved to the place location
		for i, key in enumerate(state_keys):
		if i == 0:
		next_state_example[key] = "(80,100,50)" # This is where it's placed
		else:
		next_state_example[key] = f"({i100+200},{i100+200},{50})" # Other objects remain in place

		example = f"""e.g. <think>I need to pick the {target_object.replace('_position','')} at (100,100,40) and place it at (80,100,50)</think><answer>pick(100,100,40){action_sep}place(80,100,50)</answer><next_state>{next_state_example}</next_state>"""
		return base_prompt + '\n' + example
		return base_prompt

		def grounding_worldmodeling_format_prompt(max_actions_per_step, action_sep, state_keys, add_example=True):
		state_format = {key: "(x,y,z)" for key in state_keys}

		# Create initial state example
		init_state_example = {}
		for i, key in enumerate(state_keys):
		if i == 0:
		init_state_example[key] = "(100,100,40)" # This will be the object to pick
		else:
		init_state_example[key] = f"({i100+200},{i100+200},{50})"

		# Create next state example showing the object movement
		next_state_example = {}
		for i, key in enumerate(state_keys):
		if i == 0:
		next_state_example[key] = "(80,100,50)" # This is where it's placed
		else:
		next_state_example[key] = init_state_example[key] # Other objects remain in place

		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.
		The state should be in the format of {state_format}
		Your response should be in the format of:
		<current_state>...</current_state><think>...</think><answer>...</answer><next_state>...</next_state>"""

		if add_example:
		# Use the first key as the object being manipulated in the example
		target_object = state_keys[0] if state_keys else "red_cube_position"

		example = f"""e.g. <current_state>{init_state_example}</current_state><think>I need to pick the {target_object.replace('_position','')} at (100,100,40) and place it at (80,100,50)</think><answer>pick(100,100,40){action_sep}place(80,100,50)</answer><next_state>{next_state_example}</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