arc equivalent fully observable env with 2.5x speedup

.gitignore

@@ -0,0 +1,68 @@
+*.pyc
+*.pyd
+*.pyo
+__pycache__
+*.so
+*.o
+*.c
+
+*.egg
+*.egg-info
+.*cache*/
+*venv*/
+
+/Cython/Build/*.c
+/Cython/Compiler/*.c
+/Cython/Debugger/*.c
+/Cython/Distutils/*.c
+/Cython/Parser/*.c
+/Cython/Plex/*.c
+/Cython/Runtime/refnanny.c
+/Cython/Tempita/*.c
+/Cython/*.c
+/Cython/*.html
+/Cython/*/*.html
+
+/Tools/*.elc
+/Demos/*.html
+/Demos/*/*.html
+
+/TEST_TMP/
+/build/
+/cython_build/
+cython_debug/
+/wheelhouse*/
+!tests/build/
+/dist/
+.gitrev
+.coverage
+*.patch
+*.diff
+*.orig
+*.prof
+*.rej
+*.log
+*.dep
+*.swp
+*~
+callgrind.out.*
+
+.ipynb_checkpoints
+docs/build
+docs/_build
+
+tags
+TAGS
+MANIFEST
+
+.tox
+
+# Jetbrains IDE project files
+/.idea
+/*.iml
+
+# Komodo EDIT/IDE project files
+/*.komodoproject
+
+# Visual Studio Code files
+.vscode

foarcle/__init__.py

@@ -0,0 +1,8 @@
+from .foo2arcenv import FOO2ARCv2Env
+from gymnasium.envs.registration import register
+
+register(
+    id='ARCLE/FOO2ARCv2Env-v0',
+    entry_point='foarcle.foo2arcenv:FOO2ARCv2Env',
+    max_episode_steps=1000000,
+)

foarcle/actions/color.pxd

@@ -0,0 +1,7 @@
+import numpy as np
+cimport numpy as np
+cimport cython
+
+cdef void color(unsigned char[:, ::1] grid, unsigned char[:, ::1] selection, int color) noexcept nogil
+
+cdef void floodfill(unsigned char[:, ::1] grid, tuple[int, int] grid_dim, unsigned char[:, ::1] selection, int color) noexcept nogil

foarcle/actions/color.pyx

@@ -0,0 +1,43 @@
+
+import numpy as np
+cimport numpy as np
+cimport cython
+
+@cython.boundscheck(False)
+cdef void flood_fill_recursive(unsigned char[:, ::1] grid, int grid_dim_x, int grid_dim_y, int x, int y, int color, int start_color) noexcept nogil:
+    if x < 0 or x >= grid_dim_x: return
+    if y < 0 or y >= grid_dim_y: return
+
+    if grid[x][y] == color: return
+    if grid[x][y] != start_color: return
+
+    grid[x][y] = color
+
+    flood_fill_recursive(grid, grid_dim_x, grid_dim_y, x - 1, y, color, start_color)
+    flood_fill_recursive(grid, grid_dim_x, grid_dim_y, x + 1, y, color, start_color)
+    flood_fill_recursive(grid, grid_dim_x, grid_dim_y, x, y + 1, color, start_color)
+    flood_fill_recursive(grid, grid_dim_x, grid_dim_y, x, y - 1, color, start_color)
+
+@cython.boundscheck(False)
+cdef void color(unsigned char[:, ::1] grid, unsigned char[:, ::1] selection, int color) noexcept nogil:
+    for i in range(grid.shape[0]):
+        for j in range(grid.shape[1]):
+            grid[i, j] = grid[i, j] * (1 - selection[i, j]) + color * selection[i, j]
+
+@cython.boundscheck(False)
+cdef void floodfill(unsigned char[:, ::1] grid, tuple[int, int] grid_dim, unsigned char[:, ::1] selection, int color) noexcept nogil:
+    cdef int x = -1, y = -1, 
+    cdef int i, j
+    for i in range(selection.shape[0]):
+        for j in range(selection.shape[1]):
+            if selection[i, j] == 1:
+                if  x == -1 and y == -1:
+                    x = i
+                    y = j
+                else:
+                    return
+    if x == -1 and y == -1:
+        return
+    cdef int start_color = grid[x][y]
+
+    flood_fill_recursive(grid, grid_dim[0], grid_dim[1], x, y, color, start_color)

foarcle/actions/critical.pxd

@@ -0,0 +1,5 @@
+
+import numpy as np
+cimport numpy as np
+
+cdef (int, int) crop_grid(np.ndarray[np.uint8_t, ndim=2] grid, tuple[int, int] grid_dim, np.ndarray[np.npy_bool, ndim=2, cast=True] selection)

foarcle/actions/critical.pyx

@@ -0,0 +1,28 @@
+
+import numpy as np
+cimport numpy as np
+from foarcle.actions.object cimport _get_bbox
+cdef extern from "limits.h":
+    cdef int INT_MAX
+    cdef int INT_MIN
+
+cdef (int, int) crop_grid(
+    np.ndarray[np.uint8_t, ndim=2] grid,
+    tuple[int, int] grid_dim,
+    np.ndarray[np.npy_bool, ndim=2, cast=True] selection):
+
+
+    cdef int xmin, xmax, ymin, ymax, H, W
+
+    xmin, xmax, ymin, ymax = _get_bbox(selection)
+    if xmax == INT_MIN:
+        return grid_dim
+
+    H = xmax - xmin + 1
+    W = ymax - ymin + 1
+
+    patch = np.zeros((H, W), dtype=np.uint8)
+    np.copyto(dst=patch, src=grid[xmin:xmax + 1, ymin:ymax + 1], where=selection[xmin:xmax + 1, ymin:ymax + 1])
+    grid[:, :] = 0
+    np.copyto(grid[:H, :W], patch)
+    return H, W

foarcle/actions/o2actions.pyx

@@ -0,0 +1,168 @@
+
+from foarcle.actions.color cimport color, floodfill
+from foarcle.actions.object cimport move, rotate, flip, copy, paste
+from foarcle.actions.critical cimport crop_grid
+import numpy as np
+cimport numpy as np
+cimport cython
+
+
+cpdef act(
+    np.ndarray[np.uint8_t, ndim=2] inp,
+    tuple[int, int] inp_dim,
+    np.ndarray[np.uint8_t, ndim=2] grid,
+    tuple[int, int] grid_dim,
+    np.ndarray[np.npy_bool, ndim=2, cast=True] selected,
+    np.ndarray[np.uint8_t, ndim=2] clip,
+    tuple[int, int] clip_dim, 
+    char terminated, 
+    char objsel_active, 
+    np.ndarray[np.uint8_t, ndim=2] objsel, 
+    np.ndarray[np.uint8_t, ndim=2] objsel_area, 
+    np.ndarray[np.uint8_t, ndim=2] objsel_bg, 
+    tuple[int, int] objsel_coord, 
+    char objsel_rot,
+    np.ndarray[np.npy_bool, ndim=2, cast=True] selection,
+    char operation):
+
+    if operation >= 0 and operation <= 9:
+        color(grid, selection, operation)
+        selected[:] = 0
+        objsel_active = False
+    elif operation <= 19:
+        floodfill(grid, grid_dim, selection, operation - 10)
+        selected[:] = 0
+        objsel_active = False
+    elif operation <= 23:
+        grid, selected, objsel, objsel_area, objsel_bg, objsel_coord, objsel_active, objsel_rot = move(
+            grid, grid_dim, selected, objsel, objsel_area, objsel_bg, objsel_coord, objsel_active, objsel_rot, selection, operation - 20
+        )
+    elif operation == 24:
+        grid, selected, objsel, objsel_area, objsel_bg, objsel_coord, objsel_active, objsel_rot = rotate(
+            grid, grid_dim, selected, objsel, objsel_area, objsel_bg, objsel_coord, objsel_active, objsel_rot, selection, 1
+        )
+    elif operation == 25:
+        grid, selected, objsel, objsel_area, objsel_bg, objsel_coord, objsel_active, objsel_rot = rotate(
+            grid, grid_dim, selected, objsel, objsel_area, objsel_bg, objsel_coord, objsel_active, objsel_rot, selection, 3
+        )
+    elif operation == 26:
+        grid, selected, objsel, objsel_area, objsel_bg, objsel_coord, objsel_active, objsel_rot = flip(
+            grid, grid_dim, selected, objsel, objsel_area, objsel_bg, objsel_coord, objsel_active, objsel_rot, selection, "H"
+        )
+    elif operation == 27:
+        grid, selected, objsel, objsel_area, objsel_bg, objsel_coord, objsel_active, objsel_rot = flip(
+            grid, grid_dim, selected, objsel, objsel_area, objsel_bg, objsel_coord, objsel_active, objsel_rot, selection, "V"
+        )
+    elif operation == 28:
+        clip_dim = copy(inp, inp_dim, grid, clip, clip_dim, selection, "I")
+        selected[:] = 0
+        objsel_active = False
+    elif operation == 29:
+        clip_dim = copy(inp, inp_dim, grid, clip, clip_dim, selection, "O")
+        selected[:] = 0
+        objsel_active = False
+    elif operation == 30:
+        paste(grid, clip, clip_dim, selection)
+        selected[:] = 0
+        objsel_active = False
+    elif operation == 31:
+        np.copyto(grid, inp)
+        grid_dim = inp_dim
+        selected[:] = 0
+        objsel_active = False
+    elif operation == 32:
+        grid[:] = 0
+        selected[:] = 0
+        objsel_active = False
+    elif operation == 33:
+        grid_dim = crop_grid(grid, grid_dim, selection)
+        selected[:] = 0
+        objsel_active = False
+    elif operation == 34:
+        terminated = 1
+    else:
+        assert False
+
+    return grid, grid_dim, selected, clip, clip_dim, terminated, objsel_active, objsel, objsel_area, objsel_bg, objsel_coord, objsel_rot
+
+cpdef batch_act(
+    b_inp, b_inp_dim,
+    b_grid, b_grid_dim, b_selected, b_clip, b_clip_dim, b_terminated, b_objsel_active, b_objsel, b_objsel_area, b_objsel_bg, b_objsel_coord, b_objsel_rot,
+    b_selection, b_operation):
+
+    nb_grid = b_grid.copy()
+    nb_grid_dim = b_grid_dim.copy()
+    nb_selected = b_selected.copy()
+    nb_clip = b_clip.copy()
+    nb_clip_dim = b_clip_dim.copy()
+    nb_terminated = b_terminated.copy()
+    nb_objsel_active = b_objsel_active.copy()
+    nb_objsel = b_objsel.copy()
+    nb_objsel_area = b_objsel_area.copy()
+    nb_objsel_bg = b_objsel_bg.copy()
+    nb_objsel_coord = b_objsel_coord.copy()
+    nb_objsel_rot = b_objsel_rot.copy()
+
+    for i, (
+        inp, inp_dim, grid, grid_dim, selected, clip, clip_dim, terminated, objsel_active, objsel, objsel_area, objsel_bg, objsel_coord, objsel_rot, selection, operation
+    ) in enumerate(zip(
+        b_inp,
+        b_inp_dim,
+        b_grid, 
+        b_grid_dim, 
+        b_selected, 
+        b_clip, 
+        b_clip_dim, 
+        b_terminated, 
+        b_objsel_active, 
+        b_objsel, 
+        b_objsel_area, 
+        b_objsel_bg, 
+        b_objsel_coord, 
+        b_objsel_rot, 
+        b_selection, 
+        b_operation)):
+
+        (nb_grid[i], 
+        nb_grid_dim[i], 
+        nb_selected[i], 
+        nb_clip[i], 
+        nb_clip_dim[i], 
+        nb_terminated[i], 
+        nb_objsel_active[i], 
+        nb_objsel[i], 
+        nb_objsel_area[i], 
+        nb_objsel_bg[i], 
+        nb_objsel_coord[i], 
+        nb_objsel_rot[i]) = act(
+            inp, 
+            inp_dim, 
+            grid, 
+            grid_dim, 
+            selected, 
+            clip, 
+            clip_dim, 
+            terminated, 
+            objsel_active, 
+            objsel, 
+            objsel_area, 
+            objsel_bg, 
+            objsel_coord, 
+            objsel_rot, 
+            selection, 
+            operation)
+
+    return (
+        nb_grid, 
+        nb_grid_dim, 
+        nb_selected, 
+        nb_clip, 
+        nb_clip_dim, 
+        nb_terminated, 
+        nb_objsel_active, 
+        nb_objsel, 
+        nb_objsel_area, 
+        nb_objsel_bg, 
+        nb_objsel_coord, 
+        nb_objsel_rot
+    )

foarcle/actions/object.pxd

@@ -0,0 +1,50 @@
+
+import numpy as np
+cimport numpy as np
+cimport cython
+
+cdef (int, int, int, int) _get_bbox(unsigned char[:, ::1] img) noexcept nogil
+
+cdef rotate(
+    np.ndarray[np.uint8_t, ndim=2, mode='c'] grid,
+    tuple[int, int] grid_dim,
+    np.ndarray[np.npy_bool, ndim=2, cast=True] selected,
+    np.ndarray[np.uint8_t, ndim=2, mode='c'] objsel, 
+    np.ndarray[np.uint8_t, ndim=2, mode='c'] objsel_area, 
+    np.ndarray[np.uint8_t, ndim=2] objsel_bg, 
+    tuple[int, int] objsel_coord, 
+    char objsel_active, 
+    char objsel_rot,
+    np.ndarray[np.npy_bool, ndim=2, cast=True] selection,
+    k)
+
+cdef move(
+    np.ndarray[np.uint8_t, ndim=2, mode='c'] grid,
+    tuple[int, int] grid_dim,
+    np.ndarray[np.npy_bool, ndim=2, cast=True] selected,
+    np.ndarray[np.uint8_t, ndim=2, mode='c'] objsel, 
+    np.ndarray[np.uint8_t, ndim=2, mode='c'] objsel_area, 
+    np.ndarray[np.uint8_t, ndim=2] objsel_bg, 
+    tuple[int, int] objsel_coord, 
+    char objsel_active, 
+    char objsel_rot,
+    np.ndarray[np.npy_bool, ndim=2, cast=True] selection,
+    char d)
+
+cdef flip(
+    np.ndarray[np.uint8_t, ndim=2, mode='c'] grid,
+    tuple[int, int] grid_dim,
+    np.ndarray[np.npy_bool, ndim=2, cast=True] selected,
+    np.ndarray[np.uint8_t, ndim=2, mode='c'] objsel, 
+    np.ndarray[np.uint8_t, ndim=2, mode='c'] objsel_area, 
+    np.ndarray[np.uint8_t, ndim=2] objsel_bg, 
+    tuple[int, int] objsel_coord, 
+    char objsel_active, 
+    char objsel_rot,
+    np.ndarray[np.npy_bool, ndim=2, cast=True] selection,
+    axis)
+
+
+cdef copy(np.ndarray[np.uint8_t, ndim=2] inp, inp_dim, np.ndarray[np.uint8_t, ndim=2] grid, clip, clip_dim, np.ndarray[np.npy_bool, ndim=2, cast=True] selection, source)
+
+cdef void paste(grid, np.ndarray[np.uint8_t, ndim=2] clip, clip_dim, np.ndarray[np.npy_bool, ndim=2, cast=True] selection)