synth-ai 0.2.4.dev4__py3-none-any.whl → 0.2.4.dev5__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (104) hide show
  1. synth_ai/environments/examples/__init__.py +1 -0
  2. synth_ai/environments/examples/crafter_classic/__init__.py +8 -0
  3. synth_ai/environments/examples/crafter_classic/config_logging.py +111 -0
  4. synth_ai/environments/examples/crafter_classic/debug_translation.py +0 -0
  5. synth_ai/environments/examples/crafter_classic/engine.py +575 -0
  6. synth_ai/environments/examples/crafter_classic/engine_deterministic_patch.py +63 -0
  7. synth_ai/environments/examples/crafter_classic/engine_helpers/action_map.py +5 -0
  8. synth_ai/environments/examples/crafter_classic/engine_helpers/serialization.py +74 -0
  9. synth_ai/environments/examples/crafter_classic/engine_serialization_patch_v3.py +266 -0
  10. synth_ai/environments/examples/crafter_classic/environment.py +364 -0
  11. synth_ai/environments/examples/crafter_classic/taskset.py +233 -0
  12. synth_ai/environments/examples/crafter_classic/trace_hooks_v3.py +229 -0
  13. synth_ai/environments/examples/crafter_classic/world_config_patch_simple.py +298 -0
  14. synth_ai/environments/examples/crafter_custom/__init__.py +4 -0
  15. synth_ai/environments/examples/crafter_custom/crafter/__init__.py +7 -0
  16. synth_ai/environments/examples/crafter_custom/crafter/config.py +182 -0
  17. synth_ai/environments/examples/crafter_custom/crafter/constants.py +8 -0
  18. synth_ai/environments/examples/crafter_custom/crafter/engine.py +269 -0
  19. synth_ai/environments/examples/crafter_custom/crafter/env.py +266 -0
  20. synth_ai/environments/examples/crafter_custom/crafter/objects.py +418 -0
  21. synth_ai/environments/examples/crafter_custom/crafter/recorder.py +187 -0
  22. synth_ai/environments/examples/crafter_custom/crafter/worldgen.py +119 -0
  23. synth_ai/environments/examples/crafter_custom/dataset_builder.py +373 -0
  24. synth_ai/environments/examples/crafter_custom/environment.py +312 -0
  25. synth_ai/environments/examples/crafter_custom/run_dataset.py +305 -0
  26. synth_ai/environments/examples/enron/art_helpers/email_search_tools.py +156 -0
  27. synth_ai/environments/examples/enron/art_helpers/local_email_db.py +280 -0
  28. synth_ai/environments/examples/enron/art_helpers/types_enron.py +24 -0
  29. synth_ai/environments/examples/enron/engine.py +291 -0
  30. synth_ai/environments/examples/enron/environment.py +165 -0
  31. synth_ai/environments/examples/enron/taskset.py +112 -0
  32. synth_ai/environments/examples/minigrid/__init__.py +48 -0
  33. synth_ai/environments/examples/minigrid/engine.py +589 -0
  34. synth_ai/environments/examples/minigrid/environment.py +274 -0
  35. synth_ai/environments/examples/minigrid/environment_mapping.py +242 -0
  36. synth_ai/environments/examples/minigrid/puzzle_loader.py +416 -0
  37. synth_ai/environments/examples/minigrid/taskset.py +583 -0
  38. synth_ai/environments/examples/nethack/__init__.py +7 -0
  39. synth_ai/environments/examples/nethack/achievements.py +337 -0
  40. synth_ai/environments/examples/nethack/engine.py +738 -0
  41. synth_ai/environments/examples/nethack/environment.py +255 -0
  42. synth_ai/environments/examples/nethack/helpers/__init__.py +42 -0
  43. synth_ai/environments/examples/nethack/helpers/action_mapping.py +301 -0
  44. synth_ai/environments/examples/nethack/helpers/nle_wrapper.py +401 -0
  45. synth_ai/environments/examples/nethack/helpers/observation_utils.py +433 -0
  46. synth_ai/environments/examples/nethack/helpers/recording_wrapper.py +201 -0
  47. synth_ai/environments/examples/nethack/helpers/trajectory_recorder.py +268 -0
  48. synth_ai/environments/examples/nethack/helpers/visualization/replay_viewer.py +308 -0
  49. synth_ai/environments/examples/nethack/helpers/visualization/visualizer.py +430 -0
  50. synth_ai/environments/examples/nethack/taskset.py +323 -0
  51. synth_ai/environments/examples/red/__init__.py +7 -0
  52. synth_ai/environments/examples/red/config_logging.py +110 -0
  53. synth_ai/environments/examples/red/engine.py +693 -0
  54. synth_ai/environments/examples/red/engine_helpers/__init__.py +1 -0
  55. synth_ai/environments/examples/red/engine_helpers/memory_map.py +28 -0
  56. synth_ai/environments/examples/red/engine_helpers/reward_components.py +275 -0
  57. synth_ai/environments/examples/red/engine_helpers/reward_library/__init__.py +142 -0
  58. synth_ai/environments/examples/red/engine_helpers/reward_library/adaptive_rewards.py +56 -0
  59. synth_ai/environments/examples/red/engine_helpers/reward_library/battle_rewards.py +283 -0
  60. synth_ai/environments/examples/red/engine_helpers/reward_library/composite_rewards.py +149 -0
  61. synth_ai/environments/examples/red/engine_helpers/reward_library/economy_rewards.py +137 -0
  62. synth_ai/environments/examples/red/engine_helpers/reward_library/efficiency_rewards.py +56 -0
  63. synth_ai/environments/examples/red/engine_helpers/reward_library/exploration_rewards.py +330 -0
  64. synth_ai/environments/examples/red/engine_helpers/reward_library/novelty_rewards.py +120 -0
  65. synth_ai/environments/examples/red/engine_helpers/reward_library/pallet_town_rewards.py +558 -0
  66. synth_ai/environments/examples/red/engine_helpers/reward_library/pokemon_rewards.py +312 -0
  67. synth_ai/environments/examples/red/engine_helpers/reward_library/social_rewards.py +147 -0
  68. synth_ai/environments/examples/red/engine_helpers/reward_library/story_rewards.py +246 -0
  69. synth_ai/environments/examples/red/engine_helpers/screen_analysis.py +367 -0
  70. synth_ai/environments/examples/red/engine_helpers/state_extraction.py +139 -0
  71. synth_ai/environments/examples/red/environment.py +235 -0
  72. synth_ai/environments/examples/red/taskset.py +77 -0
  73. synth_ai/environments/examples/sokoban/__init__.py +1 -0
  74. synth_ai/environments/examples/sokoban/engine.py +675 -0
  75. synth_ai/environments/examples/sokoban/engine_helpers/__init__.py +1 -0
  76. synth_ai/environments/examples/sokoban/engine_helpers/room_utils.py +656 -0
  77. synth_ai/environments/examples/sokoban/engine_helpers/vendored/__init__.py +17 -0
  78. synth_ai/environments/examples/sokoban/engine_helpers/vendored/envs/__init__.py +3 -0
  79. synth_ai/environments/examples/sokoban/engine_helpers/vendored/envs/boxoban_env.py +129 -0
  80. synth_ai/environments/examples/sokoban/engine_helpers/vendored/envs/render_utils.py +370 -0
  81. synth_ai/environments/examples/sokoban/engine_helpers/vendored/envs/room_utils.py +331 -0
  82. synth_ai/environments/examples/sokoban/engine_helpers/vendored/envs/sokoban_env.py +305 -0
  83. synth_ai/environments/examples/sokoban/engine_helpers/vendored/envs/sokoban_env_fixed_targets.py +66 -0
  84. synth_ai/environments/examples/sokoban/engine_helpers/vendored/envs/sokoban_env_pull.py +114 -0
  85. synth_ai/environments/examples/sokoban/engine_helpers/vendored/envs/sokoban_env_two_player.py +122 -0
  86. synth_ai/environments/examples/sokoban/engine_helpers/vendored/envs/sokoban_env_variations.py +394 -0
  87. synth_ai/environments/examples/sokoban/environment.py +228 -0
  88. synth_ai/environments/examples/sokoban/generate_verified_puzzles.py +438 -0
  89. synth_ai/environments/examples/sokoban/puzzle_loader.py +311 -0
  90. synth_ai/environments/examples/sokoban/taskset.py +425 -0
  91. synth_ai/environments/examples/tictactoe/__init__.py +1 -0
  92. synth_ai/environments/examples/tictactoe/engine.py +368 -0
  93. synth_ai/environments/examples/tictactoe/environment.py +239 -0
  94. synth_ai/environments/examples/tictactoe/taskset.py +214 -0
  95. synth_ai/environments/examples/verilog/__init__.py +10 -0
  96. synth_ai/environments/examples/verilog/engine.py +328 -0
  97. synth_ai/environments/examples/verilog/environment.py +349 -0
  98. synth_ai/environments/examples/verilog/taskset.py +418 -0
  99. {synth_ai-0.2.4.dev4.dist-info → synth_ai-0.2.4.dev5.dist-info}/METADATA +1 -1
  100. {synth_ai-0.2.4.dev4.dist-info → synth_ai-0.2.4.dev5.dist-info}/RECORD +104 -6
  101. {synth_ai-0.2.4.dev4.dist-info → synth_ai-0.2.4.dev5.dist-info}/WHEEL +0 -0
  102. {synth_ai-0.2.4.dev4.dist-info → synth_ai-0.2.4.dev5.dist-info}/entry_points.txt +0 -0
  103. {synth_ai-0.2.4.dev4.dist-info → synth_ai-0.2.4.dev5.dist-info}/licenses/LICENSE +0 -0
  104. {synth_ai-0.2.4.dev4.dist-info → synth_ai-0.2.4.dev5.dist-info}/top_level.txt +0 -0
@@ -0,0 +1,187 @@
1
+ import datetime
2
+ import json
3
+ import pathlib
4
+
5
+ import imageio
6
+ import numpy as np
7
+
8
+
9
+ class Recorder:
10
+ def __init__(
11
+ self,
12
+ env,
13
+ directory,
14
+ save_stats=True,
15
+ save_video=True,
16
+ save_episode=True,
17
+ video_size=(512, 512),
18
+ ):
19
+ if directory and save_stats:
20
+ env = StatsRecorder(env, directory)
21
+ if directory and save_video:
22
+ env = VideoRecorder(env, directory, video_size)
23
+ if directory and save_episode:
24
+ env = EpisodeRecorder(env, directory)
25
+ self._env = env
26
+
27
+ def __getattr__(self, name):
28
+ if name.startswith("__"):
29
+ raise AttributeError(name)
30
+ return getattr(self._env, name)
31
+
32
+
33
+ class StatsRecorder:
34
+ def __init__(self, env, directory):
35
+ self._env = env
36
+ self._directory = pathlib.Path(directory).expanduser()
37
+ self._directory.mkdir(exist_ok=True, parents=True)
38
+ self._file = (self._directory / "stats.jsonl").open("a")
39
+ self._length = None
40
+ self._reward = None
41
+ self._unlocked = None
42
+ self._stats = None
43
+
44
+ def __getattr__(self, name):
45
+ if name.startswith("__"):
46
+ raise AttributeError(name)
47
+ return getattr(self._env, name)
48
+
49
+ def reset(self):
50
+ obs = self._env.reset()
51
+ self._length = 0
52
+ self._reward = 0
53
+ self._unlocked = None
54
+ self._stats = None
55
+ return obs
56
+
57
+ def step(self, action):
58
+ obs, reward, done, info = self._env.step(action)
59
+ self._length += 1
60
+ self._reward += info["reward"]
61
+ if done:
62
+ self._stats = {"length": self._length, "reward": round(self._reward, 1)}
63
+ for key, value in info["achievements"].items():
64
+ self._stats[f"achievement_{key}"] = value
65
+ self._save()
66
+ return obs, reward, done, info
67
+
68
+ def _save(self):
69
+ self._file.write(json.dumps(self._stats) + "\n")
70
+ self._file.flush()
71
+
72
+
73
+ class VideoRecorder:
74
+ def __init__(self, env, directory, size=(512, 512)):
75
+ if not hasattr(env, "episode_name"):
76
+ env = EpisodeName(env)
77
+ self._env = env
78
+ self._directory = pathlib.Path(directory).expanduser()
79
+ self._directory.mkdir(exist_ok=True, parents=True)
80
+ self._size = size
81
+ self._frames = None
82
+
83
+ def __getattr__(self, name):
84
+ if name.startswith("__"):
85
+ raise AttributeError(name)
86
+ return getattr(self._env, name)
87
+
88
+ def reset(self):
89
+ obs = self._env.reset()
90
+ self._frames = [self._env.render(self._size)]
91
+ return obs
92
+
93
+ def step(self, action):
94
+ obs, reward, done, info = self._env.step(action)
95
+ self._frames.append(self._env.render(self._size))
96
+ if done:
97
+ self._save()
98
+ return obs, reward, done, info
99
+
100
+ def _save(self):
101
+ filename = str(self._directory / (self._env.episode_name + ".mp4"))
102
+ imageio.mimsave(filename, self._frames)
103
+
104
+
105
+ class EpisodeRecorder:
106
+ def __init__(self, env, directory):
107
+ if not hasattr(env, "episode_name"):
108
+ env = EpisodeName(env)
109
+ self._env = env
110
+ self._directory = pathlib.Path(directory).expanduser()
111
+ self._directory.mkdir(exist_ok=True, parents=True)
112
+ self._episode = None
113
+
114
+ def __getattr__(self, name):
115
+ if name.startswith("__"):
116
+ raise AttributeError(name)
117
+ return getattr(self._env, name)
118
+
119
+ def reset(self):
120
+ obs = self._env.reset()
121
+ self._episode = [{"image": obs}]
122
+ return obs
123
+
124
+ def step(self, action):
125
+ # Transitions are defined from the environment perspective, meaning that a
126
+ # transition contains the action and the resulting reward and next
127
+ # observation produced by the environment in response to said action.
128
+ obs, reward, done, info = self._env.step(action)
129
+ transition = {
130
+ "action": action,
131
+ "image": obs,
132
+ "reward": reward,
133
+ "done": done,
134
+ }
135
+ for key, value in info.items():
136
+ if key in ("inventory", "achievements"):
137
+ continue
138
+ transition[key] = value
139
+ for key, value in info["achievements"].items():
140
+ transition[f"achievement_{key}"] = value
141
+ for key, value in info["inventory"].items():
142
+ transition[f"ainventory_{key}"] = value
143
+ self._episode.append(transition)
144
+ if done:
145
+ self._save()
146
+ return obs, reward, done, info
147
+
148
+ def _save(self):
149
+ filename = str(self._directory / (self._env.episode_name + ".npz"))
150
+ # Fill in zeros for keys missing at the first time step.
151
+ for key, value in self._episode[1].items():
152
+ if key not in self._episode[0]:
153
+ self._episode[0][key] = np.zeros_like(value)
154
+ episode = {k: np.array([step[k] for step in self._episode]) for k in self._episode[0]}
155
+ np.savez_compressed(filename, **episode)
156
+
157
+
158
+ class EpisodeName:
159
+ def __init__(self, env):
160
+ self._env = env
161
+ self._timestamp = None
162
+ self._unlocked = None
163
+ self._length = None
164
+
165
+ def __getattr__(self, name):
166
+ if name.startswith("__"):
167
+ raise AttributeError(name)
168
+ return getattr(self._env, name)
169
+
170
+ def reset(self):
171
+ obs = self._env.reset()
172
+ self._timestamp = None
173
+ self._unlocked = None
174
+ self._length = 0
175
+ return obs
176
+
177
+ def step(self, action):
178
+ obs, reward, done, info = self._env.step(action)
179
+ self._length += 1
180
+ if done:
181
+ self._timestamp = datetime.datetime.now().strftime("%Y%m%dT%H%M%S")
182
+ self._unlocked = sum(int(v >= 1) for v in info["achievements"].values())
183
+ return obs, reward, done, info
184
+
185
+ @property
186
+ def episode_name(self):
187
+ return f"{self._timestamp}-ach{self._unlocked}-len{self._length}"
@@ -0,0 +1,119 @@
1
+ import functools
2
+
3
+ import numpy as np
4
+ import opensimplex
5
+
6
+ from . import constants
7
+ from . import objects
8
+ from .config import WorldGenConfig
9
+
10
+
11
+ def generate_world(world, player):
12
+ simplex = opensimplex.OpenSimplex(seed=world.random.randint(0, 2**31 - 1))
13
+ tunnels = np.zeros(world.area, bool)
14
+ for x in range(world.area[0]):
15
+ for y in range(world.area[1]):
16
+ _set_material(world, (x, y), player, tunnels, simplex)
17
+ for x in range(world.area[0]):
18
+ for y in range(world.area[1]):
19
+ _set_object(world, (x, y), player, tunnels)
20
+
21
+
22
+ def _set_material(world, pos, player, tunnels, simplex):
23
+ x, y = pos
24
+ config = getattr(world, "_config", WorldGenConfig())
25
+ simplex = functools.partial(_simplex, simplex)
26
+ uniform = world.random.uniform
27
+
28
+ # Clear area around spawn
29
+ dist_from_spawn = np.sqrt((x - player.pos[0]) ** 2 + (y - player.pos[1]) ** 2)
30
+ start = config.spawn_clear_radius - dist_from_spawn
31
+ start += 2 * simplex(x, y, 8, 3)
32
+ start = 1 / (1 + np.exp(-start))
33
+
34
+ # Terrain generation
35
+ water = simplex(x, y, 3, {15: 1, 5: 0.15}, False) + 0.1
36
+ water -= 2 * start
37
+ mountain = simplex(x, y, 0, {15: 1, 5: 0.3})
38
+ mountain -= 4 * start + 0.3 * water
39
+
40
+ if start > 0.5:
41
+ world[x, y] = "grass"
42
+ elif mountain > config.mountain_threshold:
43
+ if (simplex(x, y, 6, 7) > 0.15 and mountain > 0.3) or simplex(
44
+ x, y, 6, 5
45
+ ) > config.cave_threshold: # cave
46
+ world[x, y] = "path"
47
+ elif simplex(2 * x, y / 5, 7, 3) > 0.4: # horizonal tunnel
48
+ world[x, y] = "path"
49
+ tunnels[x, y] = True
50
+ elif simplex(x / 5, 2 * y, 7, 3) > 0.4: # vertical tunnel
51
+ world[x, y] = "path"
52
+ tunnels[x, y] = True
53
+ elif simplex(x, y, 1, 8) > config.coal_threshold and uniform() > (
54
+ 1 - config.coal_probability
55
+ ):
56
+ world[x, y] = "coal"
57
+ elif simplex(x, y, 2, 6) > config.iron_threshold and uniform() > (
58
+ 1 - config.iron_probability
59
+ ):
60
+ world[x, y] = "iron"
61
+ elif mountain > config.diamond_threshold and uniform() > (1 - config.diamond_probability):
62
+ world[x, y] = "diamond"
63
+ elif mountain > 0.3 and simplex(x, y, 6, 5) > config.lava_threshold:
64
+ world[x, y] = "lava"
65
+ else:
66
+ world[x, y] = "stone"
67
+ elif (config.sand_threshold - 0.1) < water <= config.sand_threshold and simplex(
68
+ x, y, 4, 9
69
+ ) > -0.2:
70
+ world[x, y] = "sand"
71
+ elif config.water_threshold < water:
72
+ world[x, y] = "water"
73
+ else: # grassland
74
+ if simplex(x, y, 5, 7) > config.tree_threshold and uniform() > (
75
+ 1 - config.tree_probability
76
+ ):
77
+ world[x, y] = "tree"
78
+ else:
79
+ world[x, y] = "grass"
80
+
81
+
82
+ def _set_object(world, pos, player, tunnels):
83
+ x, y = pos
84
+ config = getattr(world, "_config", WorldGenConfig())
85
+ uniform = world.random.uniform
86
+ dist = np.sqrt((x - player.pos[0]) ** 2 + (y - player.pos[1]) ** 2)
87
+ material, _ = world[x, y]
88
+
89
+ if material not in constants.walkable:
90
+ pass
91
+ elif (
92
+ dist > config.cow_min_distance
93
+ and material == "grass"
94
+ and uniform() < config.cow_spawn_probability
95
+ ):
96
+ world.add(objects.Cow(world, (x, y)))
97
+ elif (
98
+ dist > config.zombie_min_distance
99
+ and material == "grass"
100
+ and uniform() < config.zombie_spawn_probability
101
+ ):
102
+ world.add(objects.Zombie(world, (x, y), player))
103
+ elif material == "path" and tunnels[x, y] and uniform() < config.skeleton_spawn_probability:
104
+ world.add(objects.Skeleton(world, (x, y), player))
105
+
106
+
107
+ def _simplex(simplex, x, y, z, sizes, normalize=True):
108
+ if not isinstance(sizes, dict):
109
+ sizes = {sizes: 1}
110
+ value = 0
111
+ for size, weight in sizes.items():
112
+ if hasattr(simplex, "noise3d"):
113
+ noise = simplex.noise3d(x / size, y / size, z)
114
+ else:
115
+ noise = simplex.noise3(x / size, y / size, z)
116
+ value += weight * noise
117
+ if normalize:
118
+ value /= sum(sizes.values())
119
+ return value
@@ -0,0 +1,373 @@
1
+ #!/usr/bin/env python3
2
+ """
3
+ Dataset builder for Crafter Custom Environments
4
+ """
5
+
6
+ import json
7
+ import uuid
8
+ from dataclasses import dataclass, asdict
9
+ from pathlib import Path
10
+ from typing import Any, Dict, Optional, List, Set
11
+ import random
12
+ import gzip
13
+ import pickle
14
+
15
+
16
+ @dataclass
17
+ class CrafterTask:
18
+ """Defines the global task parameters for Crafter"""
19
+
20
+ global_premises: str = "You are playing Crafter, a survival game where you collect resources, craft items, and survive."
21
+ global_constraints: str = "You can only perform actions available in the action space. You must manage health and hunger."
22
+ global_objectives: str = (
23
+ "Achieve as many unique achievements as possible to maximize your score."
24
+ )
25
+ shared_env_params: Dict[str, Any] = None
26
+
27
+ def __post_init__(self):
28
+ if self.shared_env_params is None:
29
+ self.shared_env_params = {"render_mode": None, "size": (64, 64)}
30
+
31
+
32
+ @dataclass
33
+ class CrafterImpetus:
34
+ """Instructions for a specific Crafter instance"""
35
+
36
+ instructions: str
37
+ achievement_focus: Optional[List[str]] = None # Optional list of achievements to prioritize
38
+
39
+
40
+ @dataclass
41
+ class CrafterIntent:
42
+ """Success criteria and evaluation for Crafter"""
43
+
44
+ rubric: Dict[str, Any]
45
+ target_achievements: Optional[List[str]] = None
46
+ minimum_score: Optional[int] = None
47
+ gold_trajectories: Optional[List[Dict[str, Any]]] = None
48
+
49
+
50
+ @dataclass
51
+ class CrafterMetadata:
52
+ """Metadata for a Crafter instance"""
53
+
54
+ difficulty: str # easy, normal, hard, peaceful, resource_rich
55
+ world_seed: int
56
+ spawn_position: Optional[tuple] = None
57
+ initial_resources_nearby: Optional[Dict[str, int]] = None
58
+ initial_entities: Optional[Dict[str, int]] = None
59
+ config_params: Optional[Dict[str, Any]] = None
60
+
61
+
62
+ @dataclass
63
+ class SplitInfo:
64
+ """Train/val/test split information"""
65
+
66
+ val_instance_ids: Set[str]
67
+ test_instance_ids: Set[str]
68
+ _is_split_defined: bool = True
69
+
70
+
71
+ @dataclass
72
+ class CrafterInstance:
73
+ """A single Crafter environment instance"""
74
+
75
+ id: uuid.UUID
76
+ impetus: CrafterImpetus
77
+ intent: CrafterIntent
78
+ metadata: CrafterMetadata
79
+ is_reproducible: bool = True
80
+ initial_engine_snapshot: Optional[bytes] = None # Serialized engine state
81
+
82
+
83
+ @dataclass
84
+ class CrafterDataset:
85
+ """Collection of Crafter instances"""
86
+
87
+ name: str
88
+ description: str
89
+ instances: List[CrafterInstance]
90
+ split_info: SplitInfo
91
+ task: CrafterTask
92
+
93
+
94
+ class CrafterDatasetBuilder:
95
+ """Builds datasets of Crafter instances with various difficulties"""
96
+
97
+ DIFFICULTIES = ["easy", "normal", "hard", "peaceful", "resource_rich"]
98
+
99
+ # Achievement categories for focused tasks
100
+ ACHIEVEMENT_CATEGORIES = {
101
+ "basic_survival": ["collect_coal", "collect_wood", "eat_cow", "collect_sapling"],
102
+ "tools": ["make_wood_pickaxe", "make_stone_pickaxe", "make_iron_pickaxe"],
103
+ "combat": ["defeat_zombie", "defeat_skeleton", "make_wood_sword", "make_stone_sword"],
104
+ "advanced": ["make_iron_sword", "collect_diamond", "place_table", "place_plant"],
105
+ "exploration": ["collect_stone", "collect_iron", "collect_coal", "drink_water"],
106
+ }
107
+
108
+ def __init__(self, output_dir: Path = None):
109
+ self.output_dir = output_dir or Path("dataset")
110
+ self.output_dir.mkdir(exist_ok=True)
111
+
112
+ def create_instance(
113
+ self,
114
+ difficulty: str,
115
+ seed: int,
116
+ impetus_type: str = "general",
117
+ achievement_focus: Optional[List[str]] = None,
118
+ ) -> CrafterInstance:
119
+ """Create a single Crafter instance"""
120
+
121
+ # Create impetus based on type
122
+ if impetus_type == "general":
123
+ impetus = CrafterImpetus(
124
+ instructions="Survive and thrive in the Crafter world. Collect resources, craft tools, and achieve as many accomplishments as possible.",
125
+ achievement_focus=None,
126
+ )
127
+ intent = CrafterIntent(
128
+ rubric={
129
+ "description": "General survival and achievement",
130
+ "criteria": {
131
+ "achievements": "Number of unique achievements unlocked",
132
+ "survival_time": "Number of steps survived",
133
+ "resource_efficiency": "Efficient use of collected resources",
134
+ },
135
+ },
136
+ minimum_score=1 if difficulty == "easy" else (3 if difficulty == "normal" else 5),
137
+ )
138
+ elif impetus_type == "focused":
139
+ focus_category = random.choice(list(self.ACHIEVEMENT_CATEGORIES.keys()))
140
+ achievements = self.ACHIEVEMENT_CATEGORIES[focus_category]
141
+ impetus = CrafterImpetus(
142
+ instructions=f"Focus on {focus_category.replace('_', ' ')} achievements. Try to complete: {', '.join(achievements[:3])}",
143
+ achievement_focus=achievements,
144
+ )
145
+ intent = CrafterIntent(
146
+ rubric={
147
+ "description": f"Complete {focus_category} achievements",
148
+ "criteria": {
149
+ "target_achievements": f"Complete at least one of: {', '.join(achievements[:3])}",
150
+ "efficiency": "Complete target achievements quickly",
151
+ },
152
+ },
153
+ target_achievements=achievements[:3],
154
+ minimum_score=1,
155
+ )
156
+ elif impetus_type == "speedrun":
157
+ target = random.choice(["place_table", "make_iron_pickaxe", "collect_diamond"])
158
+ impetus = CrafterImpetus(
159
+ instructions=f"Speedrun challenge: {target} as quickly as possible!",
160
+ achievement_focus=[target],
161
+ )
162
+ intent = CrafterIntent(
163
+ rubric={
164
+ "description": f"Complete {target} speedrun",
165
+ "criteria": {
166
+ "completion": f"Successfully {target}",
167
+ "speed": "Complete in minimal steps",
168
+ },
169
+ },
170
+ target_achievements=[target],
171
+ minimum_score=1,
172
+ )
173
+ else:
174
+ raise ValueError(f"Unknown impetus type: {impetus_type}")
175
+
176
+ # Create metadata
177
+ metadata = CrafterMetadata(
178
+ difficulty=difficulty,
179
+ world_seed=seed,
180
+ config_params=self._get_config_params(difficulty),
181
+ )
182
+
183
+ # Create instance
184
+ instance = CrafterInstance(
185
+ id=uuid.uuid4(),
186
+ impetus=impetus,
187
+ intent=intent,
188
+ metadata=metadata,
189
+ is_reproducible=True,
190
+ initial_engine_snapshot=None,
191
+ )
192
+
193
+ return instance
194
+
195
+ def _get_config_params(self, difficulty: str) -> Dict[str, Any]:
196
+ """Get key config parameters for a difficulty"""
197
+ # These are simplified summaries - actual configs are in JSON files
198
+ params = {
199
+ "easy": {
200
+ "resource_multiplier": "~2x",
201
+ "enemy_density": "low",
202
+ "coal_probability": 0.25,
203
+ "iron_probability": 0.35,
204
+ "diamond_probability": 0.02,
205
+ },
206
+ "normal": {
207
+ "resource_multiplier": "1x",
208
+ "enemy_density": "normal",
209
+ "coal_probability": 0.15,
210
+ "iron_probability": 0.25,
211
+ "diamond_probability": 0.006,
212
+ },
213
+ "hard": {
214
+ "resource_multiplier": "~0.5x",
215
+ "enemy_density": "high",
216
+ "coal_probability": 0.08,
217
+ "iron_probability": 0.15,
218
+ "diamond_probability": 0.002,
219
+ },
220
+ "peaceful": {
221
+ "resource_multiplier": "~3x",
222
+ "enemy_density": "none",
223
+ "coal_probability": 0.30,
224
+ "iron_probability": 0.40,
225
+ "diamond_probability": 0.03,
226
+ },
227
+ "resource_rich": {
228
+ "resource_multiplier": "~8x",
229
+ "enemy_density": "minimal",
230
+ "coal_probability": 0.60,
231
+ "iron_probability": 0.70,
232
+ "diamond_probability": 0.80,
233
+ },
234
+ }
235
+ return params.get(difficulty, {})
236
+
237
+ def build_dataset(
238
+ self,
239
+ name: str,
240
+ instances_per_difficulty: Dict[str, int],
241
+ impetus_distribution: Dict[str, float] = None,
242
+ val_split: float = 0.2,
243
+ test_split: float = 0.2,
244
+ ) -> CrafterDataset:
245
+ """Build a complete dataset"""
246
+
247
+ if impetus_distribution is None:
248
+ impetus_distribution = {"general": 0.6, "focused": 0.3, "speedrun": 0.1}
249
+
250
+ instances = []
251
+
252
+ # Create instances for each difficulty
253
+ for difficulty, count in instances_per_difficulty.items():
254
+ if difficulty not in self.DIFFICULTIES:
255
+ raise ValueError(f"Unknown difficulty: {difficulty}")
256
+
257
+ for i in range(count):
258
+ # Generate unique seed
259
+ seed = random.randint(0, 1000000)
260
+
261
+ # Choose impetus type based on distribution
262
+ impetus_type = random.choices(
263
+ list(impetus_distribution.keys()), weights=list(impetus_distribution.values())
264
+ )[0]
265
+
266
+ instance = self.create_instance(
267
+ difficulty=difficulty, seed=seed, impetus_type=impetus_type
268
+ )
269
+ instances.append(instance)
270
+
271
+ # Create splits
272
+ random.shuffle(instances)
273
+ n_val = int(len(instances) * val_split)
274
+ n_test = int(len(instances) * test_split)
275
+
276
+ val_ids = {str(inst.id) for inst in instances[:n_val]}
277
+ test_ids = {str(inst.id) for inst in instances[n_val : n_val + n_test]}
278
+
279
+ split_info = SplitInfo(
280
+ val_instance_ids=val_ids, test_instance_ids=test_ids, _is_split_defined=True
281
+ )
282
+
283
+ # Create dataset
284
+ dataset = CrafterDataset(
285
+ name=name,
286
+ description=f"Crafter dataset with {len(instances)} instances across {len(instances_per_difficulty)} difficulties",
287
+ instances=instances,
288
+ split_info=split_info,
289
+ task=CrafterTask(),
290
+ )
291
+
292
+ return dataset
293
+
294
+ def save_dataset(self, dataset: CrafterDataset, format: str = "json"):
295
+ """Save dataset to disk"""
296
+ dataset_dir = self.output_dir / dataset.name
297
+ dataset_dir.mkdir(exist_ok=True)
298
+
299
+ if format == "json":
300
+ # Save metadata
301
+ metadata = {
302
+ "name": dataset.name,
303
+ "description": dataset.description,
304
+ "num_instances": len(dataset.instances),
305
+ "task": asdict(dataset.task),
306
+ "split_info": {
307
+ "val_instance_ids": list(dataset.split_info.val_instance_ids),
308
+ "test_instance_ids": list(dataset.split_info.test_instance_ids),
309
+ },
310
+ }
311
+ with open(dataset_dir / "metadata.json", "w") as f:
312
+ json.dump(metadata, f, indent=2)
313
+
314
+ # Save instances
315
+ instances_data = []
316
+ for inst in dataset.instances:
317
+ inst_data = {
318
+ "id": str(inst.id),
319
+ "impetus": asdict(inst.impetus),
320
+ "intent": asdict(inst.intent),
321
+ "metadata": asdict(inst.metadata),
322
+ "is_reproducible": inst.is_reproducible,
323
+ }
324
+ instances_data.append(inst_data)
325
+
326
+ with open(dataset_dir / "instances.json", "w") as f:
327
+ json.dump(instances_data, f, indent=2)
328
+
329
+ print(f"Dataset saved to {dataset_dir}")
330
+ print(f"Total instances: {len(dataset.instances)}")
331
+ print(f"Val instances: {len(dataset.split_info.val_instance_ids)}")
332
+ print(f"Test instances: {len(dataset.split_info.test_instance_ids)}")
333
+ print(
334
+ f"Train instances: {len(dataset.instances) - len(dataset.split_info.val_instance_ids) - len(dataset.split_info.test_instance_ids)}"
335
+ )
336
+
337
+
338
+ def main():
339
+ """Example usage"""
340
+ builder = CrafterDatasetBuilder()
341
+
342
+ # Build a balanced dataset
343
+ dataset = builder.build_dataset(
344
+ name="crafter_balanced_v1",
345
+ instances_per_difficulty={
346
+ "easy": 20,
347
+ "normal": 20,
348
+ "hard": 20,
349
+ "peaceful": 10,
350
+ "resource_rich": 10,
351
+ },
352
+ impetus_distribution={"general": 0.6, "focused": 0.3, "speedrun": 0.1},
353
+ )
354
+
355
+ builder.save_dataset(dataset)
356
+
357
+ # Build a difficulty progression dataset
358
+ dataset2 = builder.build_dataset(
359
+ name="crafter_progression_v1",
360
+ instances_per_difficulty={
361
+ "easy": 30,
362
+ "normal": 25,
363
+ "hard": 15,
364
+ "peaceful": 5,
365
+ "resource_rich": 5,
366
+ },
367
+ )
368
+
369
+ builder.save_dataset(dataset2)
370
+
371
+
372
+ if __name__ == "__main__":
373
+ main()