update frozen lake

from dataclasses import dataclass, field

from abc import ABC, abstractmethod

from typing import Optional, List, Union

@dataclass

class BaseConfig(ABC):

    format_reward: float = 0.5

    image_placeholder: str = "<image>"

    special_token_list: Optional[List[str]] = field(default_factory=lambda: ["<think>", "</think>", "<answer>", "</answer>"])

    action_sep: str = ","

    @abstractmethod

    def config_id(self) -> str: # config identifier, wandb and mllm rollout manager use this to identify the config

        pass

from vagen.env_new.base_config import BaseConfig

from dataclasses import dataclass, fields,field

from typing import Optional, List, Union

@dataclass

class FrozenLakeConfig(BaseConfig):

    desc: Optional[List[str]] = None  # environment map

    is_slippery: bool = False

    size: int = 4

    p: float = 0.8  # probability of frozen tile

    render_mode: str = "vision"  # "text" or "vision"

    max_actions_per_step: int = 3

    min_actions_to_succeed: int = 5

    def config_id(self) -> str:

        id_fields=["is_slippery", "size", "p", "render_mode", "max_actions_per_step", "min_actions_to_succeed"]

        id_str = ",".join([f"{field.name}={getattr(self, field.name)}" for field in fields(self) if field.name in id_fields])

        return f"FrozenLakeConfig({id_str})"

if __name__ == "__main__":

    config = FrozenLakeConfig()

    print(config.config_id())

 No newline at end of file

system_prompt_text = """You are a FrozenLake solver.

FrozenLake Quick Guide

Goal: Reach the goal (G).

Symbols:

_ Frozen | O Hole | G Goal | P Player | X Player fell into hole | √ Player on goal

Rules:

1. Avoid falling into holes (O).

2. Frozen tiles are slippery, you may move perpendicular to your intended direction.

Actions you can take: Left, Down, Right, Up. You can take up to {max_actions_per_step} action(s) at a time.

Left: move left to the cell to the left.

Down: move down to the cell below.

Right: move right to the cell to the right.

Up: move up to the cell above.

Rewards:

Fall into hole: 0

Reach goal: +10.0

Format correct: +0.5

Please think step by step and provide the actions you want to take.

Your response should be in the format of <think>...</think><answer>...</answer>

"""

system_prompt_vision = """You are a FrozenLake solver.

FrozenLake Quick Guide

Goal: Reach the goal (G).

Symbols:

Light blue: Frozen surface | Black: Hole | Green: Goal | Red: Player

Rules:

1. Avoid falling into holes.

2. Frozen tiles are slippery, you may move perpendicular to your intended direction.

Actions you can take: Left, Down, Right, Up. You can take up to {max_actions_per_step} action(s) at a time.

Rewards:

Fall into hole: 0

Reach goal: +10.0

Format correct: +0.5

Please think step by step and provide the actions you want to take.

Your response should be in the format of <think>...</think><answer>...</answer>

"""

init_observation_template = """

[Initial Observation]:

{observation}

Decide your next action(s).

Please think step by step and provide the actions you want to take.

Your response should be in the format of <think>...</think><answer>...</answer>

"""

action_template = """After your answer, the extracted valid action is {valid_action}.

After that, the observation is:

{observation}

reward: {reward}

done: {done}

Decide your next action(s).

Please think step by step and provide the actions you want to take.

Your response should be in the format of <think>...</think><answer>...</answer>

"""

 No newline at end of file

from typing import Dict, List, Optional, Tuple, Any

from gymnasium.utils import seeding

import numpy as np

def generate_random_map(size: int = 8, p: float = 0.8, seed: Optional[int] = None) -> List[str]:

    """Generates a random valid map (one that has a path from start to goal)

    Args:

        size: size of each side of the grid

        p: probability that a tile is frozen

        seed: optional seed to ensure the generation of reproducible maps

    Returns:

        A random valid map

    """

    valid = False

    board = []  # initialize to make pyright happy

    np_random, _ = seeding.np_random(seed)

    # generate random start and end points

    while not valid:

        p = min(1, p)

        board = np_random.choice(["F", "H"], (size, size), p=[p, 1 - p])

        while True:

            start_r = np_random.integers(0, size)

            start_c = np_random.integers(0, size)

            goal_r = np_random.integers(0, size)

            goal_c = np_random.integers(0, size)

            # Ensure start and goal are different positions

            if (start_r, start_c) != (goal_r, goal_c):

                break

        board[start_r][start_c] = "S"

        board[goal_r][goal_c] = "G"

        valid = is_valid(board, size)

    return ["".join(x) for x in board]

def is_valid(board: List[List[str]], max_size: int) -> bool:

    """Check if the board is valid (has a path from start to goal)"""

    frontier, discovered = [], set()

    # find the start point

    start_r, start_c = np.where(np.array(board) == "S")

    frontier.append((start_r[0], start_c[0]))

    # dfs to check if there is a path from start to goal

    while frontier:

        r, c = frontier.pop()

        if not (r, c) in discovered:

            discovered.add((r, c))

            directions = [(1, 0), (0, 1), (-1, 0), (0, -1)]

            for x, y in directions:

                r_new = r + x

                c_new = c + y

                if r_new < 0 or r_new >= max_size or c_new < 0 or c_new >= max_size:

                    continue

                if board[r_new][c_new] == "G":

                    return True

                if board[r_new][c_new] != "H":

                    frontier.append((r_new, c_new))

    return False

 No newline at end of file

    render_mode: str = "vision"

    min_actions_to_succeed: int = 5

    max_actions_per_step: int = 3

    format_reward = 0.5

    def config_id(self) -> str:

        return str(self)

        id_fields = ["dim_room", "max_steps", "num_boxes", "render_mode", "min_actions_to_succeed", "max_actions_per_step"]

        id_str = ",".join([f"{field.name}={getattr(self, field.name)}" for field in field(self) if field.name in id_fields])

        return f"SokobanConfig({id_str})"

if __name__ == "__main__":

    config = SokobanConfig()

    print(config)

    print(config.config_id())

 No newline at end of file