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

        with open(self.config.alf_config_path) as reader:

            alf_config = yaml.safe_load(reader)

        # @TODO Force TextWorld environment type

        if self.config.render_mode == "vision":

            alf_config['env']['type'] = 'AlfredThorEnv'

            env = alfworld.agents.environment.AlfredThorEnv(alf_config)

        else:

            alf_config['env']['type'] = 'AlfredTWEnv'

            env = alfworld.agents.environment.AlfredTWEnv(alf_config)

        # Initialize TextWorld environment

        env = alfworld.agents.environment.AlfredTWEnv(alf_config, train_eval='train')

        self.env = env.init_env(batch_size=1)

        # Track state

        self.valid_actions = []

    def step(self, llm_raw_response):

        """Process LLM response and take a step in the environment

        """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,

            max_actions=self.config.max_actions_per_step

        )

        # Extract and process action

        # Extract actions and process them

        action_list = parsed['actions']

        legal_action = False

        for i in range(len(action_list)):

            action_list[i] = action_list[i].lower()

            if len(action_list[i]) == 0:

                print("Action is empty!!!!")

                # If action is empty, choose a random action from the action list

                action_list[i] = self.prev_admissible_commands[random.randint(0, len(self.prev_admissible_commands)-1)]

            else:

                action_index = action_list[i].find('"action":')

                if action_index == -1:

                    string = action_list[i][-30:]

                else:

                    string = action_list[i][action_index:]

                for act in self.prev_admissible_commands:

                    if act in string:

                        action_list[i] = act

                        legal_action = True

                        break

                # If not a valid action, randomly pick an action

                if not legal_action:

                    action_list[i] = self.prev_admissible_commands[random.randint(0, len(self.prev_admissible_commands)-1)]

        # Use the first valid action from the action list

        action_text = action_list[0] if action_list else ""

        # Store valid action for observation formatting

        # Take the step in ALFWorld env

        obs, reward, done, infos = self.env.step([action_text])

        # Render the environment and track the action effectiveness

        observation = self._render(obs, infos)  # Add render here to capture observation

        # 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

            "llm_response": parsed

        }

        # Update total reward

        # Compute reward with a penalty for illegal actions

        if isinstance(reward, tuple):

            reward_value = reward[0]

        elif isinstance(reward, (list, np.ndarray)):

            reward_value = reward[0]

        else:

            reward_value = reward

        if reward_value is None:

            reward_value = 0  # Handle None rewards

        # Add penalty if the action is illegal

        if not legal_action:

            reward_value -= 1  # Apply penalty for illegal action

        self.total_reward += reward_value

        # Update admissible commands for next step

            )

        # For text mode, just return the observation string

        if self.config.render_mode == "text":

            return {

                "obs_str": obs_str

            }

        # @TODO

        else:

        if self.config.render_mode == "vision":

            img = self.env.get_frames()[0]

            img_placeholder = self.config.image_placeholder

            return {

                "obs_str": obs_str,

                "multi_modal_data": {

                }

                "multi_modal_data": {img_placeholder: [convert_numpy_to_PIL(img)]}

            }

        else:

            return {"obs_str": obs_str}

 No newline at end of file

#!/usr/bin/env python3

"""

start_x11.py

A self‑contained script to launch a virtual X11 server using NVIDIA and Xorg,

enabling AI2-THOR to run headlessly on your remote server.

Usage:

    python start_x11.py [DISPLAY]

    # e.g. python start_x11.py 0  <-- will start X on :0 and export DISPLAY

You can also import and call start_x11.start(display) from Python before

initializing your ALFWorldEnv or AI2-THOR Controller.

"""

import os

import re

import shlex

import platform

import subprocess

import tempfile

# =============================================================================

# Utilities to suppress ALSA errors by writing a null asound configuration

# =============================================================================

def _setup_null_asoundrc():

    """

    Create a ~/.asoundrc file that directs ALSA to the null device,

    preventing 'cannot find card' errors.

    """

    home = os.path.expanduser('~')

    cfg_path = os.path.join(home, '.asoundrc')

    if os.path.exists(cfg_path):

        return

    null_cfg = (

        "pcm.!default {\n"

        "    type null\n"

        "}\n"

        "ctl.!default {\n"

        "    type null\n"

        "}\n"

    )

    try:

        with open(cfg_path, 'w') as f:

            f.write(null_cfg)

    except Exception:

        # If writing fails (e.g., permission), silently ignore

        pass

# =============================================================================

# Core X11 startup logic

# =============================================================================

def pci_records():

    """Parse `lspci -vmm` output into list of dicts."""

    out = subprocess.check_output(shlex.split('lspci -vmm')).decode()

    recs = []

    for device in out.strip().split("\n\n"):

        d = {}

        for row in device.split("\n"):

            key, val = row.split("\t")

            d[key.split(':')[0]] = val

        recs.append(d)

    return recs

def generate_xorg_conf(bus_ids, width=1280, height=1024):

    """Generate an Xorg config that uses NVIDIA GPUs headlessly."""

    device_tpl = '''Section "Device"

    Identifier  "Device{idx}"

    Driver      "nvidia"

    VendorName  "NVIDIA Corporation"

    BusID       "{bus}"

EndSection

'''

    screen_tpl = '''Section "Screen"

    Identifier  "Screen{idx}"

    Device      "Device{idx}"

    DefaultDepth 24

    Option      "AllowEmptyInitialConfiguration" "True"

    SubSection "Display"

        Depth 24

        Virtual {w} {h}

    EndSubSection

EndSection

'''

    parts = []

    layout_lines = []

    for i, bus in enumerate(bus_ids):

        parts.append(device_tpl.format(idx=i, bus=bus))

        parts.append(screen_tpl.format(idx=i, w=width, h=height))

        layout_lines.append(f"    Screen {i} \"Screen{i}\" 0 0")

    layout = (

        "Section \"ServerLayout\"\n"

        "    Identifier \"Layout0\"\n" +

        "\n".join(layout_lines) +

        "\nEndSection\n"

    )

    return "\n".join(parts) + "\n" + layout

def start(display=0, width=1280, height=1024):

    """Launch a headless X server on the given DISPLAY index."""

    if platform.system() != 'Linux':

        raise RuntimeError("start_x11 only supports Linux")

    # suppress ALSA errors

    os.environ['SDL_AUDIODRIVER'] = 'dummy'

    _setup_null_asoundrc()

    # find NVIDIA GPUs

    buses = []

    for r in pci_records():

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

            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))

    if not buses:

        raise RuntimeError("No NVIDIA GPU found for Xorg virtual display")

    # write temporary xorg.conf

    fd, path = tempfile.mkstemp(suffix='.conf')

    conf = generate_xorg_conf(buses, width, height)

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

        f.write(conf)

    # launch Xorg silently

    cmd = (

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

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

    )

    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}"

if __name__ == '__main__':

    import sys

    disp = int(sys.argv[1]) if len(sys.argv) > 1 else 0

    start(disp)

    print(f"Use DISPLAY=:{disp} for your AI2-THOR processes.")

import os

import sys

import unittest

from dataclasses import asdict

# Add parent directory to path if needed

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from vagen.env.alfworld.env import ALFWorldEnv

from vagen.env.alfworld.env_config import ALFWorldEnvConfig

def test_alfworld_env():

    """Basic test for ALFWorldEnv functionality."""

    print("=== Testing ALFWorldEnv ===")

    # Create configuration with path to your config file

    config = ALFWorldEnvConfig(

        alf_config_path="/workspace/VAGEN/vagen/env/alfworld/data/alf-config.py",  # Update with actual path

        max_actions_per_step=1,

        action_only_prompt=False,

        render_mode="text"  # Start with text mode for simplicity

    )

    try:

        # Test initialization

        print("1. Testing initialization...")

        env = ALFWorldEnv(config)

        print("✓ Initialization successful")

        # Test system prompt

        print("\n2. Testing system prompt...")

        prompt = env.system_prompt()

        print(f"System prompt (first 100 chars): {prompt[:100]}...")

        print("✓ System prompt generated")

        # Test reset

        print("\n3. Testing environment reset...")

        obs, info = env.reset()

        print("✓ Reset successful")

        print(f"Observation keys: {obs.keys()}")

        print(f"Observation preview: {obs['obs_str'][:100]}...")

        # Test available actions

        print("\n4. Available actions:")

        if env.prev_admissible_commands:

            for i, cmd in enumerate(env.prev_admissible_commands):

                print(f"  {i+1}. {cmd}")

        # Test step with a valid action

        print("\n5. Testing step with first available action...")

        if env.prev_admissible_commands and len(env.prev_admissible_commands) > 0:

            # Use the first available action

            action = env.prev_admissible_commands[0]

            action_json = '{"action": "' + action + '"}'

            print(f"Taking action: {action}")

            next_obs, reward, done, step_info = env.step(action_json)

            print(f"Reward: {reward}")

            print(f"Done: {done}")

            print(f"Action valid: {step_info['metrics']['turn_metrics']['action_is_valid']}")

            print(f"Action effective: {step_info['metrics']['turn_metrics']['action_is_effective']}")

            print(f"New observation preview: {next_obs['obs_str'][:100]}...")

            print("✓ Step successful")

        else:

            print("No admissible commands available to test step")

        # Test compute_reward

        print("\n6. Testing compute_reward...")

        total_reward = env.compute_reward()

        print(f"Total reward: {total_reward}")

        print("✓ Compute reward successful")

        # Test close

        print("\n7. Testing environment close...")

        env.close()

        print("✓ Close successful")

        print("\n=== All tests passed! ===")

    except Exception as e:

        print(f"Error: {e}")

        import traceback

        traceback.print_exc()

if __name__ == "__main__":

    test_alfworld_env()

 No newline at end of file

    trainer.critic_warmup=0 \

    trainer.logger=['console','wandb'] \

    trainer.project_name='vagen_debug' \

    trainer.experiment_name='grpo_mask_loss_frozenlake_vision_debug' \

    trainer.experiment_name='grpo_mask_loss_alfworld_text_debug' \

    trainer.n_gpus_per_node=1 \

    trainer.nnodes=1 \

    trainer.save_freq=100 \