alfworld update

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

    # find NVIDIA GPUs

    buses = []

    for r in pci_records():

<<<<<<< HEAD

        if r.get('Vendor') == 'NVIDIA Corporation' and (

            r.get('Class','').startswith('VGA') or 

            r.get('Class','').startswith('3D')

        ):

=======

        if r.get('Vendor') == 'NVIDIA Corporation' and r.get('Class','').startswith('VGA'):

>>>>>>> 6b8f828 (alfworld update)

            slot = r['Slot']  # e.g. '01:00.0'

            parts = re.split(r'[:\.]', slot)

            buses.append('PCI:' + ':'.join(str(int(x,16)) for x in parts))

    with os.fdopen(fd, 'w') as f:

        f.write(conf)

<<<<<<< HEAD

    # launch Xorg in the foreground

=======

    # launch Xorg silently

>>>>>>> 6b8f828 (alfworld update)

    cmd = (

        f"Xorg -noreset +extension GLX +extension RANDR +extension RENDER "

        f"-config {path} :{display}"

    )

<<<<<<< HEAD

    process = subprocess.Popen(shlex.split(cmd), stdout=subprocess.PIPE, stderr=subprocess.PIPE)

    print(f"Started Xorg on DISPLAY=:{display}")

    # export DISPLAY for this process

    os.environ['DISPLAY'] = f":{display}"

    print(f"Xorg is running on DISPLAY=:{display}. You can stop it by killing the process.")

=======

    subprocess.Popen(shlex.split(cmd), stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)

    print(f"Started Xorg on DISPLAY=:{display}")

    # export DISPLAY for this process

    os.environ['DISPLAY'] = f":{display}"

>>>>>>> 6b8f828 (alfworld update)

if __name__ == '__main__':

    import sys