wehub-resource-sync
140 lines
5.4 KiB
Python
140 lines
5.4 KiB
Python
# 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 skillopt.envs.alfworld.vendor.env_base import EnvironmentManagerBase, to_numpy
|
|
from skillopt.envs.alfworld.vendor.alfworld_prompts import (
|
|
ALFWORLD_TEMPLATE,
|
|
ALFWORLD_TEMPLATE_NO_HIS,
|
|
ALFWORLD_TEMPLATE_WITH_MEMORY,
|
|
)
|
|
from skillopt.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
|