mirror of
https://github.com/open-thought/reasoning-gym.git
synced 2026-04-27 17:23:19 +00:00
feat: Add new 1D ARC task generation functions for block manipulation
This commit is contained in:
Andreas Koepf (aider)
parent
9dac01fda7
commit
dc11f88c0b
1 changed files with 208 additions and 0 deletions
|
|
@ -440,3 +440,211 @@ def task_copy_block_to_dots(size: int, rng: Random) -> Optional[Dict[str, List[i
|
||||||
answer = write_block(block_start, block, answer)
|
answer = write_block(block_start, block, answer)
|
||||||
|
|
||||||
return {"input": question, "output": answer}
|
return {"input": question, "output": answer}
|
||||||
|
|
||||||
|
def task_copy_block_to_dots_colors(size: int, rng: Random) -> Optional[Dict[str, List[int]]]:
|
||||||
|
"""Generate a task where a block pattern is copied to dot positions with matching colors."""
|
||||||
|
block_size = 3 if rng.random() < 0.5 else 5
|
||||||
|
if block_size >= size:
|
||||||
|
return None
|
||||||
|
|
||||||
|
block_color = rng.randint(1, 9)
|
||||||
|
block = [block_color] * block_size
|
||||||
|
|
||||||
|
# Generate dots with minimum distance to prevent overlap
|
||||||
|
min_gap = block_size
|
||||||
|
dot_positions = []
|
||||||
|
dot_colors = []
|
||||||
|
pos = block_size + block_size//2 + 1
|
||||||
|
|
||||||
|
while pos < size - block_size:
|
||||||
|
if rng.random() < 0.5:
|
||||||
|
dot_color = rng.randint(1, 9)
|
||||||
|
dot_positions.append(pos)
|
||||||
|
dot_colors.append(dot_color)
|
||||||
|
pos += min_gap
|
||||||
|
pos += 1
|
||||||
|
|
||||||
|
if not dot_positions:
|
||||||
|
return None
|
||||||
|
|
||||||
|
question = gen_field(size)
|
||||||
|
question = write_block(0, block, question)
|
||||||
|
for i, pos in enumerate(dot_positions):
|
||||||
|
question[pos] = dot_colors[i]
|
||||||
|
|
||||||
|
answer = gen_field(size)
|
||||||
|
answer = write_block(0, block, answer)
|
||||||
|
for i, pos in enumerate(dot_positions):
|
||||||
|
block_start = pos - block_size//2
|
||||||
|
colored_block = [dot_colors[i]] * block_size
|
||||||
|
answer = write_block(block_start, colored_block, answer)
|
||||||
|
|
||||||
|
return {"input": question, "output": answer}
|
||||||
|
|
||||||
|
def task_paint_biggest_block(size: int, rng: Random) -> Optional[Dict[str, List[int]]]:
|
||||||
|
"""Generate a task where the largest block is painted a different color."""
|
||||||
|
target_color = 1
|
||||||
|
initial_color = rng.randint(2, 9)
|
||||||
|
|
||||||
|
# Generate random blocks
|
||||||
|
question = gen_field(size)
|
||||||
|
blocks = []
|
||||||
|
pos = 0
|
||||||
|
|
||||||
|
while pos < size:
|
||||||
|
if rng.random() < 0.4 and size - pos >= 2:
|
||||||
|
block_size = rng.randint(2, min(size - pos, 6))
|
||||||
|
blocks.append((pos, block_size))
|
||||||
|
for i in range(block_size):
|
||||||
|
question[pos + i] = initial_color
|
||||||
|
pos += block_size + 1
|
||||||
|
else:
|
||||||
|
pos += 1
|
||||||
|
|
||||||
|
if len(blocks) < 2:
|
||||||
|
return None
|
||||||
|
|
||||||
|
# Find biggest block
|
||||||
|
biggest_pos, biggest_size = max(blocks, key=lambda x: x[1])
|
||||||
|
|
||||||
|
# Check if there are multiple blocks of the same size
|
||||||
|
biggest_count = sum(1 for _, size in blocks if size == biggest_size)
|
||||||
|
if biggest_count > 1:
|
||||||
|
return None
|
||||||
|
|
||||||
|
answer = question.copy()
|
||||||
|
for i in range(biggest_size):
|
||||||
|
answer[biggest_pos + i] = target_color
|
||||||
|
|
||||||
|
return {"input": question, "output": answer}
|
||||||
|
|
||||||
|
def task_sort_blocks_by_size(size: int, rng: Random) -> Optional[Dict[str, List[int]]]:
|
||||||
|
"""Generate a task where blocks are sorted by size with 1 pixel gaps."""
|
||||||
|
color = rng.randint(1, 9)
|
||||||
|
blocks = []
|
||||||
|
pos = 0
|
||||||
|
|
||||||
|
# Generate random blocks with random sizes
|
||||||
|
while pos < size:
|
||||||
|
if rng.random() < 0.4 and size - pos >= 2:
|
||||||
|
block_size = rng.randint(1, min(size - pos, 6))
|
||||||
|
blocks.append((pos, block_size))
|
||||||
|
pos += block_size + rng.randint(1, 4) # Random gaps
|
||||||
|
else:
|
||||||
|
pos += 1
|
||||||
|
|
||||||
|
if len(blocks) < 2:
|
||||||
|
return None
|
||||||
|
|
||||||
|
# Create input field
|
||||||
|
question = gen_field(size)
|
||||||
|
for pos, block_size in blocks:
|
||||||
|
for i in range(block_size):
|
||||||
|
question[pos + i] = color
|
||||||
|
|
||||||
|
# Sort blocks by size
|
||||||
|
blocks.sort(key=lambda x: x[1])
|
||||||
|
|
||||||
|
# Check if sorted blocks fit with gaps
|
||||||
|
total_space = sum(size for _, size in blocks) + len(blocks) - 1
|
||||||
|
if total_space > size:
|
||||||
|
return None
|
||||||
|
|
||||||
|
# Create answer field with sorted blocks
|
||||||
|
answer = gen_field(size)
|
||||||
|
current_pos = 0
|
||||||
|
|
||||||
|
for _, block_size in blocks:
|
||||||
|
for i in range(block_size):
|
||||||
|
answer[current_pos + i] = color
|
||||||
|
current_pos += block_size + 1 # One pixel gap
|
||||||
|
|
||||||
|
return {"input": question, "output": answer}
|
||||||
|
|
||||||
|
def task_sort_complete_sequence(size: int, rng: Random) -> Optional[Dict[str, List[int]]]:
|
||||||
|
"""Generate a task where a complete sequence of block sizes is sorted."""
|
||||||
|
# Calculate max possible block size given total array size
|
||||||
|
max_size = 1
|
||||||
|
total_space = 0
|
||||||
|
while total_space + max_size + 1 <= size:
|
||||||
|
total_space += max_size + 1
|
||||||
|
max_size += 1
|
||||||
|
max_size -= 1
|
||||||
|
|
||||||
|
if max_size < 2:
|
||||||
|
return None
|
||||||
|
|
||||||
|
color = rng.randint(1, 9)
|
||||||
|
|
||||||
|
# Create sequence of all sizes from 1 to max_size
|
||||||
|
blocks = list(range(1, max_size + 1))
|
||||||
|
rng.shuffle(blocks)
|
||||||
|
|
||||||
|
# Create input field with shuffled blocks
|
||||||
|
question = gen_field(size)
|
||||||
|
pos = 0
|
||||||
|
for block_size in blocks:
|
||||||
|
for i in range(block_size):
|
||||||
|
question[pos + i] = color
|
||||||
|
pos += block_size + 1
|
||||||
|
|
||||||
|
# Create answer field with sorted blocks
|
||||||
|
answer = gen_field(size)
|
||||||
|
pos = 0
|
||||||
|
for block_size in range(1, max_size + 1):
|
||||||
|
for i in range(block_size):
|
||||||
|
answer[pos + i] = color
|
||||||
|
pos += block_size + 1
|
||||||
|
|
||||||
|
return {"input": question, "output": answer}
|
||||||
|
|
||||||
|
def task_recolor_blocks_by_size(size: int, rng: Random) -> Optional[Dict[str, List[int]]]:
|
||||||
|
"""Generate a task where two blocks are recolored based on their size."""
|
||||||
|
# Generate two different random sizes
|
||||||
|
size1 = rng.randint(2, 8)
|
||||||
|
size2 = rng.randint(2, 8)
|
||||||
|
while size2 == size1:
|
||||||
|
size2 = rng.randint(2, 8)
|
||||||
|
|
||||||
|
# Ensure both blocks fit with at least 1 gap
|
||||||
|
if size1 + size2 + 1 > size:
|
||||||
|
return None
|
||||||
|
|
||||||
|
# Place blocks with gap
|
||||||
|
pos1 = rng.randint(0, size - (size1 + size2 + 1))
|
||||||
|
pos2 = rng.randint(pos1 + size1 + 1, size - size2)
|
||||||
|
|
||||||
|
# Create input field with both blocks color 3
|
||||||
|
question = gen_field(size)
|
||||||
|
for i in range(size1):
|
||||||
|
question[pos1 + i] = 3
|
||||||
|
for i in range(size2):
|
||||||
|
question[pos2 + i] = 3
|
||||||
|
|
||||||
|
# Create answer field with recolored blocks
|
||||||
|
answer = question.copy()
|
||||||
|
if size1 > size2:
|
||||||
|
for i in range(size1):
|
||||||
|
answer[pos1 + i] = 1
|
||||||
|
for i in range(size2):
|
||||||
|
answer[pos2 + i] = 2
|
||||||
|
else:
|
||||||
|
for i in range(size1):
|
||||||
|
answer[pos1 + i] = 2
|
||||||
|
for i in range(size2):
|
||||||
|
answer[pos2 + i] = 1
|
||||||
|
|
||||||
|
return {"input": question, "output": answer}
|
||||||
|
|
||||||
|
def task_gravity_one_step(size: int, rng: Random) -> Optional[Dict[str, List[int]]]:
|
||||||
|
"""Generate a task where non-zero elements move one step left if possible."""
|
||||||
|
question = [rng.randint(1, 9) if rng.random() < 0.5 else 0 for _ in range(size)]
|
||||||
|
answer = question.copy()
|
||||||
|
|
||||||
|
# Move each non-zero pixel one step left if possible
|
||||||
|
for i in range(1, size):
|
||||||
|
if answer[i] != 0 and answer[i-1] == 0:
|
||||||
|
answer[i-1] = answer[i]
|
||||||
|
answer[i] = 0
|
||||||
|
|
||||||
|
return {"input": question, "output": answer}
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue