primitive skill multi env training tested (ebce6a88) · Commits · 张泽凯 / VAGEN

vagen/env/primitive_skill/env.py

+0 −1

Original line number	Diff line number	Diff line
		@@ -106,7 +106,6 @@ class PrimitiveSkillEnv(BaseEnv):


		def _render(self,info,init_obs=False,valid_actions=None):
		print("info",info)
		new_info=handel_info(info.copy())
		object_positions=new_info['obj_positions']
		other_information=new_info['other_info']

vagen/env/primitive_skill/service.py

+13 −36

Original line number	Diff line number	Diff line
		@@ -6,12 +6,14 @@ from vagen.server.serial import serialize_observation
		from .env import PrimitiveSkillEnv
		from .env_config import PrimitiveSkillEnvConfig
		from ..base.base_service_config import BaseServiceConfig

		import threading
		from mani_skill.utils.building.articulations.partnet_mobility import _load_partnet_mobility_dataset, PARTNET_MOBILITY
		class PrimitiveSkillService(BaseService):
		"""
		Service class for PrimitiveSkill environments.
		Implements batch operations with parallel processing for efficiency.
		"""
		_dataset_lock = threading.Lock()

		def __init__(self, config: BaseServiceConfig):
		"""
		@@ -25,6 +27,14 @@ class PrimitiveSkillService(BaseService):
		self.environments = {}
		self.env_configs = {}

		def _ensure_dataset_initialized(self):
		"""
		Ensure the PartNet Mobility dataset is initialized in a thread-safe manner.
		"""
		with self._dataset_lock:
		if PARTNET_MOBILITY is None or "model_urdf_paths" not in PARTNET_MOBILITY:
		_load_partnet_mobility_dataset()

		def create_environments_batch(self, ids2configs: Dict[Any, Any]) -> None:
		"""
		Create multiple PrimitiveSkill environments in parallel.
		@@ -37,14 +47,14 @@ class PrimitiveSkillService(BaseService):
		- env_config: PrimitiveSkill specific configuration
		"""
		# Define worker function
		self._ensure_dataset_initialized()
		def create_single_env(env_id, config):
		# Verify environment type

		env_name = config.get('env_name', 'primitive_skill')
		if env_name != 'primitive_skill':
		return env_id, None, f"Expected environment type 'primitive_skill', got '{env_name}'"

		try:
		# Get PrimitiveSkill specific configuration
		env_config_dict = config.get('env_config', {})

		# Create environment config
		@@ -54,8 +64,6 @@ class PrimitiveSkillService(BaseService):
		env = PrimitiveSkillEnv(env_config)

		return env_id, (env, env_config), None
		except Exception as e:
		return env_id, None, str(e)

		# Use ThreadPoolExecutor for parallel creation
		with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
		@@ -92,16 +100,10 @@ class PrimitiveSkillService(BaseService):

		# Define worker function
		def reset_single_env(env_id, seed):
		try:
		if env_id not in self.environments:
		return env_id, None, f"Environment {env_id} not found"

		env = self.environments[env_id]
		observation, info = env.reset(seed=seed)
		serialized_observation = serialize_observation(observation)
		return env_id, (serialized_observation, info), None
		except Exception as e:
		return env_id, None, str(e)

		# Use ThreadPoolExecutor for parallel reset
		with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
		@@ -138,16 +140,10 @@ class PrimitiveSkillService(BaseService):

		# Define worker function
		def step_single_env(env_id, action):
		try:
		if env_id not in self.environments:
		return env_id, None, f"Environment {env_id} not found"

		env = self.environments[env_id]
		observation, reward, done, info = env.step(action)
		serialized_observation = serialize_observation(observation)
		return env_id, (serialized_observation, reward, done, info), None
		except Exception as e:
		return env_id, None, str(e)

		# Use ThreadPoolExecutor for parallel step
		with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
		@@ -183,14 +179,8 @@ class PrimitiveSkillService(BaseService):

		# Define worker function
		def compute_reward_single_env(env_id):
		try:
		if env_id not in self.environments:
		return env_id, None, f"Environment {env_id} not found"

		env = self.environments[env_id]
		return env_id, env.compute_reward(), None
		except Exception as e:
		return env_id, None, str(e)

		# Use ThreadPoolExecutor for parallel computation
		with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
		@@ -226,14 +216,8 @@ class PrimitiveSkillService(BaseService):

		# Define worker function
		def get_system_prompt_single_env(env_id):
		try:
		if env_id not in self.environments:
		return env_id, None, f"Environment {env_id} not found"

		env = self.environments[env_id]
		return env_id, env.system_prompt(), None
		except Exception as e:
		return env_id, None, str(e)

		# Use ThreadPoolExecutor for parallel retrieval
		with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
		@@ -268,15 +252,8 @@ class PrimitiveSkillService(BaseService):

		# Define worker function
		def close_single_env(env_id):
		try:
		if env_id not in self.environments:
		return f"Environment {env_id} not found"

		env = self.environments[env_id]
		env.close()
		return None
		except Exception as e:
		return str(e)

		# Use ThreadPoolExecutor for parallel closing
		with ThreadPoolExecutor(max_workers=self.max_workers) as executor:

vagen/examples/debug/debug_primitive_skill_text_grpo_service/env_config.yaml

+22 −1

Original line number	Diff line number	Diff line
		env1:
		env_name: primitive_skill
		env_config:
		render_mode: text
		env_id: "AlignTwoCube" # AlignTwoCube,PlaceTwoCube,PutAppleInDrawer,StackThreeCube
		train_size: 10000
		test_size: 8
		env2:
		env_name: primitive_skill
		env_config:
		render_mode: text
		env_id: "PlaceTwoCube" # AlignTwoCube,PlaceTwoCube,PutAppleInDrawer,StackThreeCube
		train_size: 10000
		test_size: 8
		env3:
		env_name: primitive_skill
		env_config:
		render_mode: text
		env_id: "PutAppleInDrawer" # AlignTwoCube,PlaceTwoCube,PutAppleInDrawer,StackThreeCube
		train_size: 10000
		test_size: 8
		env4:
		env_name: primitive_skill
		env_config:
		render_mode: text
		env_id: "StackThreeCube" # AlignTwoCube,PlaceTwoCube,PutAppleInDrawer,StackThreeCube
		train_size: 10000
		test_size: 32
		No newline at end of file
		test_size: 8
		No newline at end of file

vagen/mllm_agent/rollout_service.py

+4 −0

Original line number	Diff line number	Diff line
		@@ -232,8 +232,10 @@ class QwenVLRolloutManagerService():
		for bucket_key, env_ids in env_buckets.items():
		for env_id in env_ids:
		ids_to_close.append(env_id)
		self.envs.pop(env_id)

		# Step 3: Close unused environments
		#print(f"[DEBUG] ids_to_close: {ids_to_close}")
		self.env_client.close_batch(ids_to_close)
		# Step 4: Create new environments
		ids2configs_create = {}
		@@ -246,8 +248,10 @@ class QwenVLRolloutManagerService():
		ids2configs_create[id_str] = cfg
		ids2seeds_reset[id_str] = cfg["seed"]
		self.envs[id_str] = REGISTERED_ENV[cfg["env_name"]]["config_cls"](**cfg["env_config"])
		#print(f"[DEBUG] ids2configs_create: {ids2configs_create}")
		self.env_client.create_environments_batch(ids2configs_create)
		# Step 5: Reset environments
		#print(f"[DEBUG] ids2seeds_reset: {ids2seeds_reset}")
		reset_results=self.env_client.reset_batch(ids2seeds_reset)

vagen/server/server.yaml

+1 −1

Original line number	Diff line number	Diff line
		@@ -15,5 +15,5 @@ navigation:
		max_workers: 10
		devices: [0,1,2,3]
		primitive_skill:
		max_workers: 10
		max_workers: 20

Admin message