init alfworld (a3007feb) · Commits · 张泽凯 / VAGEN

vagen/env/init.py

+10 −2

Original line number	Diff line number	Diff line
		@@ -2,6 +2,10 @@ from .sokoban import SokobanEnv,SokobanEnvConfig
		from .frozenlake import FrozenLakeEnv,FrozenLakeEnvConfig, FrozenLakeService
		# from .navigation import NavigationEnv, NavigationEnvConfig, NavigationServiceConfig, NavigationService
		# from .svg import SVGEnv, SvgEnvConfig, SVGService, SVGServiceConfig
		# from .primitive_skill import PrimitiveSkillEnv, PrimitiveSkillEnvConfig, PrimitiveSkillService, PrimitiveSkillConfig
		from .alfworld import ALFWorldEnv, ALFWorldEnvConfig
		# from .navigation import NavigationEnv, NavigationEnvConfig, NavigationServiceConfig, NavigationService
		# from .svg import SVGEnv, SvgEnvConfig, SVGService, SVGServiceConfig
		# from .primitive_skill import PrimitiveSkillEnv, PrimitiveSkillEnvConfig, PrimitiveSkillService, PrimitiveSkillServiceConfig

		REGISTERED_ENV = {
		@@ -30,6 +34,10 @@ REGISTERED_ENV = {
		# "env_cls": PrimitiveSkillEnv,
		# "config_cls": PrimitiveSkillEnvConfig,
		# "service_cls": PrimitiveSkillService,
		# "service_config_cls": PrimitiveSkillServiceConfig
		# }
		# "service_config_cls": PrimitiveSkillConfig
		# },
		"alfworld": {
		"env_cls": ALFWorldEnv,
		"config_cls": ALFWorldEnvConfig,
		},
		}
		No newline at end of file

vagen/env/alfworld/init.py

0 → 100644

+3 −0

Original line number	Diff line number	Diff line
		from .env import ALFWorldEnv
		from .env_config import ALFWorldEnvConfig

vagen/env/alfworld/alf_utils.py

0 → 100644

+146 −0

Original line number	Diff line number	Diff line
		import os
		import yaml
		import torchvision.transforms as T
		from alfworld.agents.environment.alfred_thor_env import AlfredThorEnv
		import gymnasium as gym
		from gymnasium import spaces
		import alfworld.agents.environment as environment
		from typing import Optional
		import numpy as np
		import torch
		import random


		ALF_ACTION_LIST=["pass", "goto", "pick", "put", "open", "close", "toggle", "heat", "clean", "cool", "slice", "inventory", "examine", "look"]
		# ALF_ITEM_LIST =

		def load_config_file(path):
		print(f"[DEBUG]: alfconfig path: {path}")
		assert os.path.exists(path), "Invalid config file"
		with open(path) as reader:
		config = yaml.safe_load(reader)
		return config

		def get_obs_image(env):
		transform = T.Compose([T.ToTensor()])
		current_frames = env.get_frames()
		image_tensors = [transform(i).cuda() for i in current_frames]
		for i in range(len(image_tensors)):
		image_tensors[i] = image_tensors[i].permute(1, 2, 0)
		image_tensors[i]*= 255
		image_tensors[i] = image_tensors[i].int()
		image_tensors[i] = image_tensors[i][:,:,[2,1,0]]
		image_tensors = torch.stack(image_tensors, dim=0)
		return image_tensors

		class AlfEnv(gym.Env):
		def __init__(self, config_file):
		config = load_config_file(config_file)
		env_type = config['env']['type']
		env = getattr(environment, env_type)(config, train_eval='train')
		self.env = env.init_env(batch_size=1)
		self.action_space = spaces.Discrete(len(ALF_ACTION_LIST))
		self.observation_space = spaces.Box(low=0, high=255, shape=(300, 300, 3), dtype=np.uint8)
		# Add the previous admissible commands for step
		self.prev_admissible_commands = None
		self.num_envs = 1
		def step(self, action):
		## SZ.3.4: sanity checking legal action as rewards
		action, legal_action = process_action(self.env, action, self.prev_admissible_commands)
		obs, scores, dones, infos = self.env.step(action)
		infos['observation_text'] = obs
		reward = compute_reward(infos, legal_action)
		self.prev_admissible_commands = list(infos['admissible_commands'])[0]
		return self._get_obs(), reward, dones, infos

		def reset(
		self,
		seed=42,
		):
		self.env.seed(seed)
		obs, infos = self.env.reset()
		infos['observation_text'] = obs
		self.prev_admissible_commands = list(infos['admissible_commands'])[0]
		return self._get_obs(), infos

		def _get_obs(self):
		image = get_obs_image(self.env)
		return image

		def process_action(env, action=None, action_list=None):
		"""
		An function to process the action
		env: the environment should be of type AlfredThorEnv
		action: the list of action to be processeed, it is a list of strings.
		"""
		if type(env) != AlfEnv and type(env) != AlfredThorEnv:
		pass
		else:
		legal_action = False
		for i in range(len(action)):
		action[i] = action[i].lower()
		# TODO: need to figure this out
		if len(action[i]) == 0:
		print("Action is empty!!!!")
		# randomly choose an action from the action list if illegal
		action[i] = action_list[random.randint(0, len(action_list)-1)]
		else:
		try:
		action_index = action[i].find('"action":')
		# string has the following format '"action": "look"\n}'
		if action_index == -1:
		# if we cannot find "action":, then we pick the last 30 characters
		string = action[i][-30:]
		else:
		string = action[i][action_index:]
		# post processing by removing the first and last part of the string
		for act in action_list:
		if act in string:
		action[i] = act
		# if found legal action, set legal_action = True
		legal_action = True
		break
		except:
		# randomly choose an action from the action list if illegal
		action[i] = action_list[random.randint(0, len(action_list)-1)]

		return action, legal_action


		def compute_reward(infos, legal_action):
		# A function to compute the shaped reward for the alfworld environment
		# infos: the info returned by the environment
		# legal_action: a boolean value to indicate if the action is legal
		## Tentative rewards: r = success_reward * 10 + goal_conditioned_r - 1*illegal_action
		reward = 50*float(infos['won'][0]) + float(infos['goal_condition_success_rate'][0])
		if not legal_action:
		# adding a reward penalty to illegal actions
		reward -= 1
		reward = [reward]
		return torch.tensor(reward)

		def get_encoded_text(observation_text, tokenizer, model):

		encoded_input = tokenizer(observation_text, return_tensors='pt')
		outputs = model(**encoded_input)
		cls_embeddings = outputs.last_hidden_state[:,0,:]

		return cls_embeddings

		def get_concat(obs, infos, tokenizer, model, device):
		assert 'observation_text' in infos.keys(), 'observation_text not in infos!'
		obs_text = infos['observation_text']
		obs_text_encode = get_encoded_text(obs_text, tokenizer, model)
		obs_text_encode = obs_text_encode.to(device)
		obs_cat = torch.cat((obs.flatten(start_dim=1), obs_text_encode), dim=1)
		return obs_cat

		def get_cards_concat(obs, infos, tokenizer, model, device):
		## Need to move these codes to a CNN utils or something
		assert 'Formula' in infos[0].keys(), 'Formula not in infos!'
		infos = infos[0]
		formula_list = infos['Formula']
		formula = "".join([str("".join([str(x) for x in formula_list]))])
		obs_text_encode = get_encoded_text(formula, tokenizer, model).to(device)
		obs_cat = torch.cat((obs.flatten(start_dim=1), obs_text_encode), dim=1)
		return obs_cat
		No newline at end of file

vagen/env/alfworld/data/alf-config.yaml

0 → 100644

+145 −0

Original line number	Diff line number	Diff line
		dataset:
		data_path: '$ALFWORLD_DATA/json_2.1.1/train'
		eval_id_data_path: '$ALFWORLD_DATA/json_2.1.1/valid_seen' # null/None to disable
		eval_ood_data_path: '$ALFWORLD_DATA/json_2.1.1/valid_unseen' # null/None to disable
		num_train_games: 10 # max training games (<=0 indicates full dataset)
		num_eval_games: 2 # max evaluation games (<=0 indicates full dataset)

		logic:
		domain: '$ALFWORLD_DATA/logic/alfred.pddl' # PDDL domain file that defines the world dynamics
		grammar: '$ALFWORLD_DATA/logic/alfred.twl2' # Grammar file that defines the text feedbacks

		env:
		type: 'AlfredThorEnv' # 'AlfredTWEnv' or 'AlfredThorEnv' or 'AlfredHybrid'
		regen_game_files: False # check if game is solvable by expert and save to game.tw-pddl file
		domain_randomization: False # shuffle Textworld print order and object id nums
		task_types: [1, 2, 3, 4, 5, 6] # task-type ids: 1 - Pick & Place, 2 - Examine in Light, 3 - Clean & Place, 4 - Heat & Place, 5 - Cool & Place, 6 - Pick Two & Place
		expert_timeout_steps: 150 # max steps before timeout for expert to solve the task
		expert_type: "handcoded" # 'handcoded' or 'planner'. Note: the planner is very slow for real-time use
		goal_desc_human_anns_prob: 0.0 # prob of using human-annotated goal language instead of templated goals (1.0 indicates all human annotations from ALFRED)

		hybrid:
		start_eps: 100000 # starting episode of hybrid training, tw-only training upto this point
		thor_prob: 0.5 # prob of AlfredThorEnv during hybrid training
		eval_mode: "thor" # 'tw' or 'thor' - env used for evaluation during hybrid training

		thor:
		screen_width: 300 # width of THOR window
		screen_height: 300 # height of THOR window
		smooth_nav: False # smooth rotations, looks, and translations during navigation (very slow)
		save_frames_to_disk: False # save frame PNGs to disk (useful for making videos)
		save_frames_path: './videos/' # path to save frame PNGs

		controller:
		type: 'oracle' # 'oracle' or 'oracle_astar' or 'mrcnn' or 'mrcnn_astar' (aka BUTLER)
		debug: False
		load_receps: True # load receptacle locations from precomputed dict (if available)

		mask_rcnn:
		pretrained_model_path: '$ALFWORLD_DATA/detectors/mrcnn.pth'

		general:
		random_seed: 42
		use_cuda: True # disable this when running on machine without cuda
		visdom: False # plot training/eval curves, run with visdom server
		task: 'alfred'
		training_method: 'dagger' # 'dqn' or 'dagger'
		save_path: './training/' # path to save pytorch models
		observation_pool_capacity: 3 # k-size queue, 0 indicates no observation
		hide_init_receptacles: False # remove initial observation containing navigable receptacles

		training:
		batch_size: 1
		max_episode: 50000
		smoothing_eps: 0.1
		optimizer:
		learning_rate: 0.001
		clip_grad_norm: 5

		evaluate:
		run_eval: True
		batch_size: 1
		env:
		type: "AlfredTWEnv"

		checkpoint:
		report_frequency: 1000 # report every N episode
		experiment_tag: 'test' # name of experiment
		load_pretrained: False # during test, enable this so that the agent load your pretrained model
		load_from_tag: 'not loading anything' # name of pre-trained model to load in save_path

		model:
		encoder_layers: 1
		decoder_layers: 1
		encoder_conv_num: 5
		block_hidden_dim: 64
		n_heads: 1
		dropout: 0.1
		block_dropout: 0.1
		recurrent: True

		rl:
		action_space: "admissible" # 'admissible' (candidates from text engine) or 'generation' (seq2seq-style generation) or 'beam_search_choice' or 'exhaustive' (not working)
		max_target_length: 20 # max token length for seq2seq generation
		beam_width: 10 # 1 means greedy
		generate_top_k: 3

		training:
		max_nb_steps_per_episode: 50 # terminate after this many steps
		learn_start_from_this_episode: 0 # delay updates until this epsiode
		target_net_update_frequency: 500 # sync target net with online net per this many epochs

		replay:
		accumulate_reward_from_final: True
		count_reward_lambda: 0.0 # 0 to disable
		novel_object_reward_lambda: 0.0 # 0 to disable
		discount_gamma_game_reward: 0.9
		discount_gamma_count_reward: 0.5
		discount_gamma_novel_object_reward: 0.5
		replay_memory_capacity: 500000 # adjust this depending on your RAM size
		replay_memory_priority_fraction: 0.5
		update_per_k_game_steps: 5
		replay_batch_size: 64
		multi_step: 3
		replay_sample_history_length: 4
		replay_sample_update_from: 2

		epsilon_greedy:
		noisy_net: False # if this is true, then epsilon greedy is disabled
		epsilon_anneal_episodes: 1000 # -1 if not annealing
		epsilon_anneal_from: 0.3
		epsilon_anneal_to: 0.1

		dagger:
		action_space: "generation" # 'admissible' (candidates from text engine) or 'generation' (seq2seq-style generation) or 'exhaustive' (not working)
		max_target_length: 20 # max token length for seq2seq generation
		beam_width: 10 # 1 means greedy
		generate_top_k: 5
		unstick_by_beam_search: False # use beam-search for failed actions, set True during evaluation

		training:
		max_nb_steps_per_episode: 50 # terminate after this many steps

		fraction_assist:
		fraction_assist_anneal_episodes: 50000
		fraction_assist_anneal_from: 1.0
		fraction_assist_anneal_to: 0.01

		fraction_random:
		fraction_random_anneal_episodes: 0
		fraction_random_anneal_from: 0.0
		fraction_random_anneal_to: 0.0

		replay:
		replay_memory_capacity: 500000
		update_per_k_game_steps: 5
		replay_batch_size: 64
		replay_sample_history_length: 4
		replay_sample_update_from: 2

		vision_dagger:
		model_type: "resnet" # 'resnet' (whole image features) or 'maskrcnn_whole' (whole image MaskRCNN feats) or 'maskrcnn' (top k MaskRCNN detection feats) or 'no_vision' (zero vision input)
		resnet_fc_dim: 64
		maskrcnn_top_k_boxes: 10 # top k box features
		use_exploration_frame_feats: False # append feats from initial exploration (memory intensive!)
		sequence_aggregation_method: "average" # 'sum' or 'average' or 'rnn'
		No newline at end of file

vagen/env/alfworld/env.py

0 → 100644

+233 −0

Original line number	Diff line number	Diff line
		from vagen.env.base.base_env import BaseEnv
		from vagen.env.utils.context_utils import parse_llm_raw_response, convert_numpy_to_PIL
		from alfworld.agents.utils.misc import get_templated_task_desc
		from .env_config import ALFWorldEnvConfig
		from .prompt import system_prompt_text, system_prompt_vision, init_observation_template, action_template

		import alfworld.agents.environment
		import numpy as np
		import torch

		class ALFWorldEnv(BaseEnv):
		"""ALFWorld environment adapter that maps the BaseEnv interface to ALFWorld interface"""

		def __init__(self, config: ALFWorldEnvConfig):
		"""Initialize the ALFWorld environment"""
		super().__init__()
		self.config = config

		# Load ALFWorld config
		import yaml
		with open(self.config.alf_config_path) as reader:
		alf_config = yaml.safe_load(reader)

		# @TODO Force TextWorld environment type
		alf_config['env']['type'] = 'AlfredTWEnv'

		# Initialize TextWorld environment
		env = alfworld.agents.environment.AlfredTWEnv(alf_config, train_eval='train')
		self.env = env.init_env(batch_size=1)

		# Track state
		self.total_reward = 0
		self.prev_admissible_commands = None
		self.valid_actions = []

		def step(self, llm_raw_response):
		"""Process LLM response and take a step in the environment

		Args:
		llm_raw_response: Raw text response from LLM

		Returns:
		Observation, reward, done, info
		"""
		# Parse LLM response
		parsed = parse_llm_raw_response(
		response=llm_raw_response,
		special_token_list=self.config.special_token_list,
		action_sep=self.config.action_sep,
		max_actions=self.config.max_actions_per_step
		)

		# Extract and process action
		action_list = parsed['actions']
		action_text = action_list[0] if action_list else ""

		# Store valid action for observation formatting
		self.valid_actions = [action_text] if action_text else []

		# Check if action is valid
		action_is_valid = action_text in self.prev_admissible_commands

		# Take the step in ALFWorld env
		obs, reward, done, infos = self.env.step([action_text])

		# Simple tracking of state change (text environments don't have position)
		action_is_effective = len(obs[0]) > 10 # Basic check if we got a meaningful observation

		# Build observation and info
		observation = self._render(obs, infos)

		# Check if metrics are available in infos
		# Some environments might not provide all metrics
		success = False
		goal_condition_rate = 0.0

		if 'won' in infos:
		success = float(infos['won'][0]) if isinstance(infos['won'], (list, tuple)) else float(infos['won'])

		if 'goal_condition_success_rate' in infos:
		goal_condition_rate = float(infos['goal_condition_success_rate'][0]) if isinstance(infos['goal_condition_success_rate'], (list, tuple)) else float(infos['goal_condition_success_rate'])

		metrics = {
		"turn_metrics": {
		"action_is_valid": action_is_valid,
		"action_is_effective": action_is_effective,
		},
		"traj_metrics": {
		"success": success,
		"goal_condition_success_rate": goal_condition_rate
		},
		}

		info = {
		"metrics": metrics,
		"llm_raw_response": llm_raw_response,
		"llm_response": parsed
		}

		# Update total reward
		if isinstance(reward, tuple):
		reward_value = reward[0]
		elif isinstance(reward, (list, np.ndarray)):
		reward_value = reward[0]
		else:
		reward_value = reward
		self.total_reward += reward_value

		# Update admissible commands for next step
		self.prev_admissible_commands = infos['admissible_commands'][0]

		# Convert done to boolean if it's a list or array
		done_value = done[0] if isinstance(done, (list, np.ndarray)) else done

		return observation, reward_value, done_value, info

		def reset(self, seed=None):
		"""Reset the environment

		Args:
		seed: Random seed for reproducibility

		Returns:
		Observation dict, info dict
		"""
		# Handle seed manually if provided @TODO figure out better random way
		if seed is not None:
		import random
		random.seed(seed)

		np.random.seed(seed)

		if torch:
		torch.manual_seed(seed)
		if torch.cuda.is_available():
		torch.cuda.manual_seed(seed)

		obs, infos = self.env.reset()
		self.total_reward = 0
		self.prev_admissible_commands = infos['admissible_commands'][0]
		self.valid_actions = []
		return self._render(obs, infos, init_obs=True), infos

		def system_prompt(self):
		"""Generate system prompt

		Returns:
		System prompt string
		"""
		if self.config.render_mode == "vision":
		return system_prompt_vision.format(
		max_actions_per_step=self.config.max_actions_per_step,
		action_sep=self.config.action_sep
		)
		else:
		return system_prompt_text.format(
		max_actions_per_step=self.config.max_actions_per_step,
		action_sep=self.config.action_sep
		)

		def compute_reward(self):
		"""Return total reward

		Returns:
		Total reward for the episode
		"""
		return self.total_reward

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

		def _render(self, obs, infos, init_obs=False):
		"""Render the environment as observation

		This method creates a text representation of the environment state.
		In the future, it could be extended to support visual rendering.

		Args:
		obs: Raw observations from ALFWorld
		infos: Additional information from environment
		init_obs: Whether this is the initial observation

		Returns:
		Dict: Observation dictionary
		"""
		# Get the observation text
		observation_text = obs[0]

		# Format the list of admissible commands
		commands_text = "\n".join([f"'{s}'" for s in self.prev_admissible_commands]) if self.prev_admissible_commands else ""

		# Select appropriate template based on whether this is initial observation
		if init_obs:
		obs_str = init_observation_template.format(
		observation=observation_text,
		commands=commands_text
		)
		else:
		# For non-initial observations, include action results
		obs_str = action_template.format(
		valid_action=self.valid_actions[0] if self.valid_actions else "None",
		observation=observation_text,
		commands=commands_text,
		reward=self.total_reward
		)

		# Add response format instructions based on config
		if not self.config.action_only_prompt:
		obs_str += (
		"\nYour response should be a valid JSON: \n{\n"
		"\"thoughts\": \"your reasoning\", \n"
		"\"action\": \"chosen_action\"\n}"
		)
		else:
		obs_str += (
		"\nYour response should be a valid JSON: \n{\n"
		"\"action\": \"chosen_action\"\n}"
		)

		# For text mode, just return the observation string
		if self.config.render_mode == "text":
		return {
		"obs_str": obs_str
		}
		# @TODO
		else:
		img_placeholder = self.config.image_placeholder
		return {
		"obs_str": obs_str,
		"multi_modal_data": {
		}
		}
		No newline at end of file

Admin message