Initial commit
This commit is contained in:
carpedkm
+9
@@ -0,0 +1,9 @@
|
||||
"""Vendored ALFWorld environment runtime.
|
||||
|
||||
Minimal subset of SkillRL's agent_system package needed to run
|
||||
ALFWorld environments with ReflACT. Original source:
|
||||
https://github.com/NTU-LANTERN/SkillRL (Apache-2.0 License)
|
||||
"""
|
||||
from .alfworld_envs import AlfworldEnvs, build_alfworld_envs
|
||||
from .alfworld_projection import alfworld_projection
|
||||
from .env_manager import AlfWorldEnvironmentManager
|
||||
+221
@@ -0,0 +1,221 @@
|
||||
# Vendored from SkillRL (Apache-2.0 License)
|
||||
# Original: agent_system/environments/env_package/alfworld/envs.py
|
||||
# Modified: imports use pip-installed alfworld package instead of vendored copy.
|
||||
|
||||
import os
|
||||
import multiprocessing as mp
|
||||
import traceback
|
||||
import yaml
|
||||
import gymnasium as gym
|
||||
import numpy as np
|
||||
|
||||
from alfworld.agents.environment import get_environment
|
||||
|
||||
|
||||
def load_config_file(path):
|
||||
assert os.path.exists(path), f"Invalid config file: {path}"
|
||||
with open(path) as reader:
|
||||
config = yaml.safe_load(reader)
|
||||
return config
|
||||
|
||||
|
||||
def compute_reward(info, multi_modal=False):
|
||||
if multi_modal:
|
||||
reward = 10.0 * float(info['won']) + float(info['goal_condition_success_rate'])
|
||||
else:
|
||||
reward = 10.0 * float(info['won'])
|
||||
return reward
|
||||
|
||||
|
||||
class AlfworldWorker:
|
||||
"""Stateful worker that holds one ALFWorld sub-environment."""
|
||||
|
||||
def __init__(self, config, seed, base_env, gamefile=None):
|
||||
if gamefile:
|
||||
base_env.game_files = [gamefile]
|
||||
if hasattr(base_env, "num_games"):
|
||||
base_env.num_games = 1
|
||||
self.env = base_env.init_env(batch_size=1)
|
||||
self.env.seed(seed)
|
||||
|
||||
def step(self, action):
|
||||
actions = [action]
|
||||
obs, scores, dones, infos = self.env.step(actions)
|
||||
infos['observation_text'] = obs
|
||||
return obs, scores, dones, infos
|
||||
|
||||
def reset(self):
|
||||
obs, infos = self.env.reset()
|
||||
infos['observation_text'] = obs
|
||||
return obs, infos
|
||||
|
||||
|
||||
def _worker_loop(cmd_q, result_q, config, seed, is_train, eval_dataset, gamefile):
|
||||
"""Run one ALFWorld environment in a child process."""
|
||||
try:
|
||||
env_type = config['env']['type']
|
||||
base_env = get_environment(env_type)(
|
||||
config,
|
||||
train_eval='train' if is_train else eval_dataset,
|
||||
)
|
||||
worker = AlfworldWorker(config, seed, base_env, gamefile)
|
||||
result_q.put((True, "ready"))
|
||||
except BaseException:
|
||||
result_q.put((False, traceback.format_exc()))
|
||||
return
|
||||
|
||||
while True:
|
||||
cmd, payload = cmd_q.get()
|
||||
if cmd == "close":
|
||||
result_q.put((True, None))
|
||||
return
|
||||
try:
|
||||
if cmd == "reset":
|
||||
result = worker.reset()
|
||||
elif cmd == "step":
|
||||
result = worker.step(payload)
|
||||
else:
|
||||
raise ValueError(f"Unknown ALFWorld worker command: {cmd}")
|
||||
result_q.put((True, result))
|
||||
except BaseException:
|
||||
result_q.put((False, traceback.format_exc()))
|
||||
|
||||
|
||||
class _ProcessWorker:
|
||||
"""Small stdlib actor wrapper for one environment process."""
|
||||
|
||||
def __init__(self, ctx, config, seed, is_train, eval_dataset, gamefile=None):
|
||||
self.cmd_q = ctx.Queue(maxsize=1)
|
||||
self.result_q = ctx.Queue(maxsize=1)
|
||||
self.process = ctx.Process(
|
||||
target=_worker_loop,
|
||||
args=(self.cmd_q, self.result_q, config, seed, is_train, eval_dataset, gamefile),
|
||||
)
|
||||
self.process.start()
|
||||
ok, payload = self.result_q.get()
|
||||
if not ok:
|
||||
self.close(kill=True)
|
||||
raise RuntimeError(f"Failed to start ALFWorld worker:\n{payload}")
|
||||
|
||||
def send(self, cmd, payload=None):
|
||||
self.cmd_q.put((cmd, payload))
|
||||
|
||||
def recv(self):
|
||||
ok, payload = self.result_q.get()
|
||||
if not ok:
|
||||
raise RuntimeError(f"ALFWorld worker failed:\n{payload}")
|
||||
return payload
|
||||
|
||||
def close(self, kill=False):
|
||||
if self.process.is_alive() and not kill:
|
||||
try:
|
||||
self.send("close")
|
||||
self.recv()
|
||||
except Exception:
|
||||
kill = True
|
||||
if kill and self.process.is_alive():
|
||||
self.process.terminate()
|
||||
self.process.join(timeout=5)
|
||||
if self.process.is_alive():
|
||||
self.process.kill()
|
||||
self.process.join(timeout=1)
|
||||
self.cmd_q.close()
|
||||
self.result_q.close()
|
||||
|
||||
|
||||
class AlfworldEnvs(gym.Env):
|
||||
"""Vectorized ALFWorld environment using local process workers."""
|
||||
|
||||
def __init__(self, alf_config_path, seed, env_num, group_n,
|
||||
resources_per_worker, is_train=True, env_kwargs=None, gamefiles=None):
|
||||
super().__init__()
|
||||
if env_kwargs is None:
|
||||
env_kwargs = {}
|
||||
|
||||
eval_dataset = env_kwargs.get('eval_dataset', 'eval_in_distribution')
|
||||
config = load_config_file(alf_config_path)
|
||||
env_type = config['env']['type']
|
||||
self.multi_modal = (env_type == 'AlfredThorEnv')
|
||||
self.num_processes = env_num * group_n
|
||||
self.group_n = group_n
|
||||
self.gamefiles = list(gamefiles or [])
|
||||
if self.gamefiles and len(self.gamefiles) != self.num_processes:
|
||||
raise ValueError(
|
||||
f"Expected {self.num_processes} gamefiles, got {len(self.gamefiles)}"
|
||||
)
|
||||
|
||||
start_method = os.environ.get("ALFWORLD_WORKER_START_METHOD") or None
|
||||
ctx = mp.get_context(start_method) if start_method else mp.get_context()
|
||||
self.workers = []
|
||||
for i in range(self.num_processes):
|
||||
worker_gamefile = self.gamefiles[i] if self.gamefiles else None
|
||||
worker = _ProcessWorker(
|
||||
ctx,
|
||||
config,
|
||||
seed + (i // self.group_n),
|
||||
is_train,
|
||||
eval_dataset,
|
||||
worker_gamefile,
|
||||
)
|
||||
self.workers.append(worker)
|
||||
|
||||
self.prev_admissible_commands = [None for _ in range(self.num_processes)]
|
||||
|
||||
def step(self, actions):
|
||||
assert len(actions) == self.num_processes
|
||||
|
||||
for i, worker in enumerate(self.workers):
|
||||
worker.send("step", actions[i])
|
||||
results = [worker.recv() for worker in self.workers]
|
||||
|
||||
text_obs_list = []
|
||||
rewards_list = []
|
||||
dones_list = []
|
||||
info_list = []
|
||||
|
||||
for i, (obs, scores, dones, info) in enumerate(results):
|
||||
for k in info.keys():
|
||||
info[k] = info[k][0]
|
||||
text_obs_list.append(obs[0])
|
||||
dones_list.append(dones[0])
|
||||
info_list.append(info)
|
||||
self.prev_admissible_commands[i] = info['admissible_commands']
|
||||
rewards_list.append(compute_reward(info, self.multi_modal))
|
||||
|
||||
image_obs_list = None
|
||||
return text_obs_list, image_obs_list, rewards_list, dones_list, info_list
|
||||
|
||||
def reset(self):
|
||||
for worker in self.workers:
|
||||
worker.send("reset")
|
||||
results = [worker.recv() for worker in self.workers]
|
||||
|
||||
text_obs_list = []
|
||||
info_list = []
|
||||
|
||||
for i, (obs, info) in enumerate(results):
|
||||
for k in info.keys():
|
||||
info[k] = info[k][0]
|
||||
text_obs_list.append(obs[0])
|
||||
self.prev_admissible_commands[i] = info['admissible_commands']
|
||||
info_list.append(info)
|
||||
|
||||
image_obs_list = None
|
||||
return text_obs_list, image_obs_list, info_list
|
||||
|
||||
@property
|
||||
def get_admissible_commands(self):
|
||||
return self.prev_admissible_commands
|
||||
|
||||
def close(self):
|
||||
for worker in self.workers:
|
||||
worker.close()
|
||||
|
||||
|
||||
def build_alfworld_envs(alf_config_path, seed, env_num, group_n,
|
||||
resources_per_worker, is_train=True, env_kwargs=None, gamefiles=None):
|
||||
"""Build vectorized ALFWorld environments."""
|
||||
return AlfworldEnvs(
|
||||
alf_config_path, seed, env_num, group_n,
|
||||
resources_per_worker, is_train, env_kwargs, gamefiles,
|
||||
)
|
||||
@@ -0,0 +1,60 @@
|
||||
# Vendored from SkillRL (Apache-2.0 License)
|
||||
# Original: agent_system/environments/env_package/alfworld/projection.py
|
||||
|
||||
from typing import List
|
||||
import re
|
||||
|
||||
|
||||
def alfworld_projection(actions: List[str], action_pools: List[List[str]]):
|
||||
"""Process raw model outputs into valid ALFWorld actions.
|
||||
|
||||
Extracts text from ``<action>...</action>`` tags and validates that
|
||||
the response also contains ``<think>...</think>`` tags.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
actions : list[str]
|
||||
Raw model outputs, one per environment.
|
||||
action_pools : list[list[str]]
|
||||
Admissible action lists per environment (unused but kept for API compat).
|
||||
|
||||
Returns
|
||||
-------
|
||||
actions : list[str]
|
||||
Cleaned action strings.
|
||||
valids : list[int]
|
||||
1 if the action was successfully parsed, 0 otherwise.
|
||||
"""
|
||||
valids = [0] * len(actions)
|
||||
|
||||
for i in range(len(actions)):
|
||||
original_str = actions[i]
|
||||
actions[i] = actions[i].lower()
|
||||
|
||||
start_tag = "<action>"
|
||||
end_tag = "</action>"
|
||||
start_idx = actions[i].find(start_tag)
|
||||
end_idx = actions[i].find(end_tag)
|
||||
try:
|
||||
if start_idx == -1 or end_idx == -1:
|
||||
actions[i] = actions[i][-30:]
|
||||
continue
|
||||
|
||||
extracted_action = actions[i][start_idx + len(start_tag):end_idx].strip().lower()
|
||||
actions[i] = extracted_action
|
||||
valids[i] = 1
|
||||
|
||||
except Exception:
|
||||
actions[i] = actions[i][-30:]
|
||||
|
||||
# Require <think>...</think>
|
||||
think_start_idx = original_str.find("<think>")
|
||||
think_end_idx = original_str.find("</think>")
|
||||
if think_start_idx == -1 or think_end_idx == -1:
|
||||
valids[i] = 0
|
||||
|
||||
# Reject responses containing Chinese characters
|
||||
if re.search(r'[\u4e00-\u9fff]', original_str):
|
||||
valids[i] = 0
|
||||
|
||||
return actions, valids
|
||||
@@ -0,0 +1,8 @@
|
||||
# Vendored from SkillRL (Apache-2.0 License)
|
||||
# Original: agent_system/environments/prompts/alfworld.py
|
||||
|
||||
from reflact.prompts import load_prompt
|
||||
|
||||
ALFWORLD_TEMPLATE_NO_HIS = load_prompt("rollout_no_history", env="alfworld")
|
||||
ALFWORLD_TEMPLATE = load_prompt("rollout_with_history", env="alfworld")
|
||||
ALFWORLD_TEMPLATE_WITH_MEMORY = load_prompt("rollout_with_memory", env="alfworld")
|
||||
+145
@@ -0,0 +1,145 @@
|
||||
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: -1 # max training games (<=0 indicates full dataset)
|
||||
num_eval_games: -1 # 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: 'AlfredTWEnv' # '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: "tw" # '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: 10
|
||||
max_episode: 50000
|
||||
smoothing_eps: 0.1
|
||||
optimizer:
|
||||
learning_rate: 0.001
|
||||
clip_grad_norm: 5
|
||||
|
||||
evaluate:
|
||||
run_eval: True
|
||||
batch_size: 10
|
||||
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'
|
||||
+84
@@ -0,0 +1,84 @@
|
||||
# Vendored from SkillRL (Apache-2.0 License)
|
||||
# Original: agent_system/environments/base.py
|
||||
# Trimmed to only include what ALFWorld needs.
|
||||
|
||||
from typing import List, Tuple, Dict, Any
|
||||
import numpy as np
|
||||
from collections import defaultdict
|
||||
|
||||
|
||||
def to_numpy(data):
|
||||
"""Convert data to numpy array."""
|
||||
# Lazy-check for torch.Tensor to avoid hard dependency on torch
|
||||
_torch_tensor = None
|
||||
try:
|
||||
import torch
|
||||
_torch_tensor = torch.Tensor
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
if _torch_tensor is not None and isinstance(data, _torch_tensor):
|
||||
data = data.detach().cpu().numpy()
|
||||
elif isinstance(data, np.ndarray):
|
||||
pass
|
||||
elif isinstance(data, (int, float, bool, Tuple, List)):
|
||||
data = np.array(data)
|
||||
else:
|
||||
raise ValueError(f"Unsupported type: {type(data)})")
|
||||
return data
|
||||
|
||||
|
||||
class EnvironmentManagerBase:
|
||||
"""Base class for vectorized environment managers.
|
||||
|
||||
Manages a set of parallel environments, handles action projection,
|
||||
observation post-processing, and history tracking.
|
||||
"""
|
||||
|
||||
def __init__(self, envs, projection_f, config):
|
||||
self.envs = envs
|
||||
self.projection_f = projection_f
|
||||
self.config = config
|
||||
|
||||
def reset(self, kwargs) -> Dict[str, Any]:
|
||||
obs, infos = self.envs.reset()
|
||||
return {'text': None, 'image': obs, 'anchor': None}, infos
|
||||
|
||||
def step(self, text_actions: List[str]):
|
||||
actions, valids = self.projection_f(text_actions)
|
||||
next_obs, rewards, dones, infos = self.envs.step(actions)
|
||||
|
||||
next_observations = {
|
||||
'text': None,
|
||||
'image': next_obs,
|
||||
'anchor': None,
|
||||
}
|
||||
for i, info in enumerate(infos):
|
||||
info['is_action_valid'] = to_numpy(valids[i])
|
||||
|
||||
rewards = to_numpy(rewards)
|
||||
dones = to_numpy(dones)
|
||||
return next_observations, rewards, dones, infos
|
||||
|
||||
def close(self) -> None:
|
||||
self.envs.close()
|
||||
|
||||
def success_evaluator(self, *args, **kwargs) -> Dict[str, np.ndarray]:
|
||||
total_infos = kwargs['total_infos']
|
||||
total_batch_list = kwargs['total_batch_list']
|
||||
batch_size = len(total_batch_list)
|
||||
|
||||
success = defaultdict(list)
|
||||
for bs in range(batch_size):
|
||||
self._process_batch(bs, total_batch_list, total_infos, success)
|
||||
assert len(success['success_rate']) == batch_size
|
||||
return {key: np.array(value) for key, value in success.items()}
|
||||
|
||||
def _process_batch(self, batch_idx, total_batch_list, total_infos, success):
|
||||
for i in reversed(range(len(total_batch_list[batch_idx]))):
|
||||
batch_item = total_batch_list[batch_idx][i]
|
||||
if batch_item['active_masks']:
|
||||
info = total_infos[batch_idx][i]
|
||||
won_value = float(info['won'])
|
||||
success['success_rate'].append(won_value)
|
||||
return
|
||||
+139
@@ -0,0 +1,139 @@
|
||||
# Vendored from SkillRL (Apache-2.0 License)
|
||||
# Original: agent_system/environments/env_manager.py
|
||||
# Trimmed to only include AlfWorldEnvironmentManager and its helpers.
|
||||
|
||||
from typing import List, Dict, Any
|
||||
from collections import defaultdict
|
||||
import numpy as np
|
||||
|
||||
from reflact.envs.alfworld.vendor.env_base import EnvironmentManagerBase, to_numpy
|
||||
from reflact.envs.alfworld.vendor.alfworld_prompts import (
|
||||
ALFWORLD_TEMPLATE,
|
||||
ALFWORLD_TEMPLATE_NO_HIS,
|
||||
ALFWORLD_TEMPLATE_WITH_MEMORY,
|
||||
)
|
||||
from reflact.envs.alfworld.vendor.memory import SimpleMemory
|
||||
|
||||
|
||||
def parse_gamefile(infos):
|
||||
gamefile = []
|
||||
for info in infos:
|
||||
if 'extra.gamefile' in info:
|
||||
gamefile.append(info['extra.gamefile'])
|
||||
else:
|
||||
gamefile.append(None)
|
||||
return gamefile
|
||||
|
||||
|
||||
def set_gamefile(infos, gamefile):
|
||||
for i in range(len(infos)):
|
||||
if 'extra.gamefile' in infos[i]:
|
||||
infos[i]['extra.gamefile'] = gamefile[i]
|
||||
else:
|
||||
infos[i]['extra.gamefile'] = None
|
||||
return infos
|
||||
|
||||
|
||||
class AlfWorldEnvironmentManager(EnvironmentManagerBase):
|
||||
"""Manages parallel ALFWorld environments with observation templating."""
|
||||
|
||||
def __init__(self, envs, projection_f, config):
|
||||
self.memory = SimpleMemory()
|
||||
self.retrieval_memory = None
|
||||
super().__init__(envs, projection_f, config)
|
||||
|
||||
def reset(self, kwargs):
|
||||
text_obs, image_obs, infos = self.envs.reset()
|
||||
self.gamefile = parse_gamefile(infos)
|
||||
self.memory.reset(batch_size=len(text_obs))
|
||||
self.tasks = []
|
||||
self.pre_text_obs = text_obs
|
||||
self.extract_task(text_obs)
|
||||
|
||||
full_text_obs = self.build_text_obs(text_obs, self.envs.get_admissible_commands, init=True)
|
||||
return {'text': full_text_obs, 'image': image_obs, 'anchor': text_obs}, infos
|
||||
|
||||
def step(self, text_actions: List[str]):
|
||||
actions, valids = self.projection_f(text_actions, self.envs.get_admissible_commands)
|
||||
text_obs, image_obs, rewards, dones, infos = self.envs.step(actions)
|
||||
self.memory.store({'text_obs': self.pre_text_obs, 'action': actions})
|
||||
self.pre_text_obs = text_obs
|
||||
|
||||
full_text_obs = self.build_text_obs(text_obs, self.envs.get_admissible_commands)
|
||||
if infos[0].get("extra.gamefile") is None:
|
||||
infos = set_gamefile(infos, self.gamefile)
|
||||
|
||||
for i, info in enumerate(infos):
|
||||
info['is_action_valid'] = to_numpy(valids[i])
|
||||
|
||||
next_observations = {'text': full_text_obs, 'image': image_obs, 'anchor': text_obs}
|
||||
rewards = to_numpy(rewards)
|
||||
dones = to_numpy(dones)
|
||||
return next_observations, rewards, dones, infos
|
||||
|
||||
def extract_task(self, text_obs: List[str]):
|
||||
for obs in text_obs:
|
||||
task_start = obs.find('Your task is to: ')
|
||||
if task_start != -1:
|
||||
self.tasks.append(obs[task_start + len('Your task is to: '):].strip())
|
||||
else:
|
||||
raise ValueError("Task description not found in text observation.")
|
||||
|
||||
def build_text_obs(self, text_obs: List[str], admissible_actions: List[List[str]], init: bool = False) -> List[str]:
|
||||
postprocess_text_obs = []
|
||||
if not init and self.config.env.history_length > 0:
|
||||
memory_contexts, valid_lens = self.memory.fetch(
|
||||
self.config.env.history_length,
|
||||
obs_key="text_obs",
|
||||
action_key="action",
|
||||
)
|
||||
|
||||
for i in range(len(text_obs)):
|
||||
reformatted_admissible_actions = "\n ".join(
|
||||
f"'{s}'" for s in admissible_actions[i] if s != 'help'
|
||||
)
|
||||
|
||||
if init or self.config.env.history_length <= 0:
|
||||
obs = ALFWORLD_TEMPLATE_NO_HIS.format(
|
||||
current_observation=text_obs[i],
|
||||
admissible_actions=reformatted_admissible_actions,
|
||||
)
|
||||
else:
|
||||
obs = ALFWORLD_TEMPLATE.format(
|
||||
task_description=self.tasks[i],
|
||||
step_count=len(self.memory[i]),
|
||||
history_length=valid_lens[i],
|
||||
action_history=memory_contexts[i],
|
||||
current_step=len(self.memory[i]) + 1,
|
||||
current_observation=text_obs[i],
|
||||
admissible_actions=reformatted_admissible_actions,
|
||||
)
|
||||
postprocess_text_obs.append(obs)
|
||||
return postprocess_text_obs
|
||||
|
||||
def _process_batch(self, batch_idx, total_batch_list, total_infos, success):
|
||||
for i in reversed(range(len(total_batch_list[batch_idx]))):
|
||||
batch_item = total_batch_list[batch_idx][i]
|
||||
if batch_item['active_masks']:
|
||||
info = total_infos[batch_idx][i]
|
||||
won_value = float(info['won'])
|
||||
success['success_rate'].append(won_value)
|
||||
|
||||
gamefile = info.get("extra.gamefile")
|
||||
if gamefile:
|
||||
self._process_gamefile(gamefile, won_value, success)
|
||||
return
|
||||
|
||||
def _process_gamefile(self, gamefile, won_value, success):
|
||||
tasks = [
|
||||
"pick_and_place",
|
||||
"pick_two_obj_and_place",
|
||||
"look_at_obj_in_light",
|
||||
"pick_heat_then_place_in_recep",
|
||||
"pick_cool_then_place_in_recep",
|
||||
"pick_clean_then_place_in_recep",
|
||||
]
|
||||
for task in tasks:
|
||||
if task in gamefile:
|
||||
success[f"{task}_success_rate"].append(won_value)
|
||||
break
|
||||
+87
@@ -0,0 +1,87 @@
|
||||
# Vendored from SkillRL (Apache-2.0 License)
|
||||
# Original: agent_system/memory/base.py + agent_system/memory/memory.py
|
||||
# Merged into a single file for simplicity.
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import List, Dict, Any, Tuple
|
||||
|
||||
|
||||
class BaseMemory(ABC):
|
||||
"""Base class for memory management."""
|
||||
|
||||
@abstractmethod
|
||||
def __len__(self):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def __getitem__(self, idx: int):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def reset(self, batch_size: int):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def store(self, record: Dict[str, List[Any]]):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def fetch(self, step: int):
|
||||
pass
|
||||
|
||||
|
||||
class SimpleMemory(BaseMemory):
|
||||
"""Per-environment history buffer for storing observations and actions."""
|
||||
|
||||
def __init__(self):
|
||||
self._data = None
|
||||
self.keys = None
|
||||
self.batch_size = 0
|
||||
|
||||
def __len__(self):
|
||||
return len(self._data)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
return self._data[idx]
|
||||
|
||||
def reset(self, batch_size: int):
|
||||
if self._data is not None:
|
||||
self._data.clear()
|
||||
self._data = [[] for _ in range(batch_size)]
|
||||
self.batch_size = batch_size
|
||||
self.keys = None
|
||||
|
||||
def store(self, record: Dict[str, List[Any]]):
|
||||
if self.keys is None:
|
||||
self.keys = list(record.keys())
|
||||
assert self.keys == list(record.keys())
|
||||
|
||||
for env_idx in range(self.batch_size):
|
||||
self._data[env_idx].append({k: record[k][env_idx] for k in self.keys})
|
||||
|
||||
def fetch(
|
||||
self,
|
||||
history_length: int,
|
||||
obs_key: str = "text_obs",
|
||||
action_key: str = "action",
|
||||
) -> Tuple[List[str], List[int]]:
|
||||
memory_contexts, valid_lengths = [], []
|
||||
|
||||
for env_idx in range(self.batch_size):
|
||||
recent = self._data[env_idx][-history_length:]
|
||||
valid_len = len(recent)
|
||||
start_idx = len(self._data[env_idx]) - valid_len
|
||||
|
||||
lines = []
|
||||
for j, rec in enumerate(recent):
|
||||
step_num = start_idx + j + 1
|
||||
act = rec[action_key]
|
||||
obs = rec[obs_key]
|
||||
lines.append(
|
||||
f"[Observation {step_num}: '{obs}', Action {step_num}: '{act}']"
|
||||
)
|
||||
|
||||
memory_contexts.append("\n".join(lines))
|
||||
valid_lengths.append(valid_len)
|
||||
|
||||
return memory_contexts, valid_lengths
|
||||
Reference in New Issue
Block a user