1from typing import Dict, List, TypeVar
 2
 3from hcraft.elements import Item, Zone
 4from hcraft.task import GetItemTask
 5from hcraft.transformation import Transformation, Use, Yield, PLAYER, CURRENT_ZONE
 6
 7from hcraft.examples.minicraft.minicraft import MiniCraftEnv
 8
 9
10MINICRAFT_NAME = "FourRooms"
11__doc__ = MiniCraftEnv.description(MINICRAFT_NAME, for_module_header=True)
12
13
14class MiniHCraftFourRooms(MiniCraftEnv):
15    MINICRAFT_NAME = MINICRAFT_NAME
16    __doc__ = MiniCraftEnv.description(MINICRAFT_NAME)
17
18    SOUTH_WEST_ROOM = Zone("South-West")
19    """South west room."""
20
21    SOUTH_EAST_ROOM = Zone("South-East")
22    """South east room."""
23
24    NORTH_WEST_ROOM = Zone("North-West")
25    """North west room."""
26
27    NORTH_EAST_ROOM = Zone("North-East")
28    """North east room."""
29
30    GOAL = Item("goal")
31    """Goal to reach."""
32
33    def __init__(self, **kwargs) -> None:
34        self.task = GetItemTask(self.GOAL)
35
36        self.ROOMS = [
37            self.SOUTH_WEST_ROOM,
38            self.SOUTH_EAST_ROOM,
39            self.NORTH_EAST_ROOM,
40            self.NORTH_WEST_ROOM,
41        ]
42
43        super().__init__(
44            self.MINICRAFT_NAME,
45            start_zone=self.SOUTH_WEST_ROOM,
46            purpose=self.task,
47            **kwargs,
48        )
49
50    def build_transformations(self) -> List[Transformation]:
51        find_goal = Transformation(
52            "find-goal",
53            inventory_changes=[Yield(CURRENT_ZONE, self.GOAL, max=0)],
54            zone=self.NORTH_EAST_ROOM,
55        )
56
57        reach_goal = Transformation(
58            "reach-goal",
59            inventory_changes=[
60                Use(CURRENT_ZONE, self.GOAL, consume=1),
61                Yield(PLAYER, self.GOAL),
62            ],
63        )
64
65        neighbors = _get_rooms_connections(self.ROOMS)
66
67        moves = []
68        for destination, neighbor_rooms in neighbors.items():
69            for neighbor_room in neighbor_rooms:
70                moves.append(
71                    Transformation(
72                        f"go-to-{destination.name}-from-{neighbor_room.name}",
73                        destination=destination,
74                        zone=neighbor_room,
75                    )
76                )
77
78        return [find_goal, reach_goal] + moves
79
80
81Room = TypeVar("Room")
82
83
84def _get_rooms_connections(rooms: List[Room]) -> Dict[Room, List[Room]]:
85    neighbors = {}
86    for room_id, destination in enumerate(rooms):
87        neighbors[destination] = [rooms[room_id - 1], rooms[(room_id + 1) % len(rooms)]]
88    return neighbors