Implementing new substrate with other tensor libraries

[frotaur]

Hello,

i'd really like to try out asal. However, most of my Alife models I have are implemented in PyTorch. I thought that I could implement my substrate by simply wrapping my torch model, and doing the jax conversion in and out of the methods that are being exposed to asal.

However, this proved much more challenging than anticipated. I don't know JAX, but I keep receiving errors such as :

File "/home/frotaur/asal/main_opt.py", line 71, in calc_loss rollout_data = rollout_fn(rng, params) File "/home/frotaur/asal/rollout.py", line 84, in rollout_simulation _, state_vid = jax.lax.scan(step_fn, s0, split(rng, rollout_steps)) ~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/frotaur/asal/rollout.py", line 82, in step_fn next_state = substrate.step_state(_rng, state, params) File "/home/frotaur/asal/substrates/__init__.py", line 80, in step_state return self.substrate.step_state(rng, state, params) ~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^ File "/home/frotaur/asal/substrates/difflenia_sub.py", line 56, in step_state if(torch.allclose(torch.tensor(params['mu'])-self.lenia.params['mu'])): ~~~~~~~~~~~~^^^^^^^^^^^^^^ File "/home/frotaur/.pvenv/lib/python3.13/site-packages/torch/utils/dlpack.py", line 101, in from_dlpack device = ext_tensor.__dlpack_device__() ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ AttributeError: DynamicJaxprTracer has no attribute __dlpack_device__

which I assume are due to my torch operation not being JIT-compilable.

So I was wondering : do you think it is possible to implement the actual Alife models in something else than Jax, and use the substrate class to wrap the model, to be able to use asal?

Or are we stuck having to re-implement everything in Jax to use this codebase?

I can provide a minimal implementation of GoL in torch, as well as my attempt to create a substrate wrapper for it, if necessary.

Thanks for the help !

[akarshkumar0101]

I don't think it will be possible to use this jax codebase, since it needs to do end-to-end jitting.
However, I would be willing to create a PyTorch version of this library because it is clearly more popular!
Can you provide your GoL in torch code?

[frotaur]

Hey, sorry for the delay. I tried to do a minimal implementation with outer-holistic automata in torch. I implemented a 'substrate' in the spirit of Asal, although in torch, and added some questions/comments in the docstrings. Don't hesitate to let me know if there are problems, and if need be, I'm happy to help adding the pytorch functionality.

For some reason, I can't link '.py' files, so here is the code.

ca2d.py file to put into substrates/

import torch
import torch.nn.functional as F
import colorsys


class CA2D():
    """
    2D Cellular Automaton, with outer-holistic rules and binary values.
    The state of a pixel in the next generation depends on its own state and
    the sum of the values of its neighbors.
    """

    def __init__(self, size, s_num="23", b_num="3", random=False, device="cpu"):
        """
        Params :
        size : tuple, size of the automaton
        s_num : str, rule for survival
        b_num : str, rule for birth
        random : bool, if True, the initial state of the automaton is random. Otherwise, a small random square is placed in the middle.
        """
        self.h, self.w = size
        self.size = size

        self._worldmap = torch.zeros((3, self.h, self.w), dtype=float)  # (3,H,W), contains a 2D 'view' of the CA world

        self.s_num = self.get_num_from_rule(s_num)  # Translate string to number form
        self.b_num = self.get_num_from_rule(b_num)  # Translate string to number form
        self.random = random
        self.device = device

        self.world = torch.zeros((self.h, self.w), dtype=torch.int, device=device)
        self.reset()

        self.kernel = torch.tensor([[1, 1, 1], [1, 0, 1], [1, 1, 1]], device=device, dtype=torch.float)[
            None, None, :, :
        ]  # (1,1,3,3) as required by Pytorch

        self.change_highlight_color()
        self.decay_speed = 0.1

        self._worldmap = self._worldmap.to(device)

    def change_highlight_color(self):
        """
        Changes the highlight color, gets random hue
        """
        hue = torch.rand(1).item()
        light = 0.5
        saturation = 0.5
        self.highlight_color = torch.tensor(
            colorsys.hls_to_rgb(hue, saturation, light), dtype=torch.float, device=self.device
        )

    def get_num_from_rule(self, num : str):
        """
        Get the rule number for the automaton, given the string representation of the rule.
        """
        rule_out = 0
        if num != "":  # If the num is empty, the associated number is 0
            rule_out = sum([2 ** int(d) for d in num])

        return rule_out

    def get_rule_from_num(self, num : int):
        """
            Get the rule string from the rule number
        """
        numbers_list = [int(num//2**i % 2) for i in range(9)]

        out = ""
        for i, n in enumerate(numbers_list):
            if n == 1:
                out += str(i)

        return out

    def reset(self):
        """
        Resets the automaton to the initial state.
        """
        self._worldmap = torch.zeros((3, self.h, self.w), device=self.device)

        if self.random:
            self.world = self.get_init_mat(0.5)
        else:
            self.world = torch.zeros_like(self.world, dtype=torch.int, device=self.device)
            self.world[self.h // 2 - 1 : self.h // 2 + 1, self.w // 2 - 1 : self.w // 2 + 1] = torch.randint(
                0, 2, (2, 2), device=self.device
            )

    def pure_render_state(self, state):
        """
        Returns the current state of the automaton without any modifications.
        """
        return state[None, :, :].expand(3, -1, -1).to(dtype=torch.float).to(self.device)

    def draw(self):
        """
        Updates the worldmap with the current state of the automaton.
        """
        echo = torch.clamp(self._worldmap - self.decay_speed * self.highlight_color[:, None, None], min=0, max=1).to(
            self.device
        )
        self._worldmap = torch.clamp(
            self.world[None, :, :].expand(3, -1, -1).to(dtype=torch.float) + echo, min=0, max=1
        ).to(self.device)

    def change_num(self, s_num: int | str, b_num: int | str):
        """
        	Changes the rule of the automaton to the one specified by s_num and b_num
        """
        if(isinstance(s_num, str)):
            s_num = self.get_num_from_rule(s_num)
        if(isinstance(b_num, str)):
            b_num = self.get_num_from_rule(b_num)
            
        self.s_num = s_num
        self.b_num = b_num

    def step(self):
        # Generate tensors for all 8 neighbors
        w, e = self.world.roll(-1, 0), self.world.roll(1, 0)
        n, s = self.world.roll(-1, 1), self.world.roll(1, 1)
        sw, se = w.roll(1, 1), e.roll(1, 1)
        nw, ne = w.roll(-1, 1), e.roll(-1, 1)

        count = w + e + n + s + sw + se + nw + ne

        # count2 = self.convolve_count(self.world)
        # assert torch.all(count==count2)

        self.world = torch.where(
            self.world == 1, self.get_nth_bit(self.s_num, count), self.get_nth_bit(self.b_num, count)
        ).to(torch.int)

    def convolve_count(self, world):
        """
        Convolves the world with the kernel to get the count of neighbors
        """
        pad_world = F.pad(
            world[None, None, :, :].to(torch.float), (1, 1, 1, 1), mode="circular"
        )  # (1,1,H+2,W+2) circular padded

        return F.conv2d(pad_world, self.kernel)[0, 0, :, :].to(torch.int)  # (H,W)

    def get_nth_bit(self, num, s):
        """
        Get the nth bit of the number num
        """
        return (num >> s) & 1

    def get_random_rule(self, generator=None):
        """
            Return random outer holistic rules with S and B rules as integers.
        """
        b_rule = torch.randint(0, 2**9, (1,),generator=generator).item()
        s_rule = torch.randint(0, 2**9, (1,),generator=generator).item()

        return {'b_num': self.get_rule_from_num(b_rule), 's_num': self.get_rule_from_num(s_rule)}
    
    def get_init_mat(self, rand_portion):
        """
        Get initialization matrix for CA

        Params :
        rand_portion : float, portion of the screen filled with noise.
        """
        batched_size = torch.tensor([self.h, self.w])
        randsize = (batched_size * rand_portion).to(dtype=torch.int16)  # size of the random square
        randstarts = (batched_size * (1 - rand_portion) / 2).to(
            dtype=torch.int16
        )  # Where to start the index for the random square

        randsquare = torch.where(torch.randn(*randsize.tolist()) > 0, 1, 0)  # Creates random square of 0s and 1s

        init_mat = torch.zeros((self.h, self.w), dtype=torch.int16)
        init_mat[randstarts[0] : randstarts[0] + randsize[0], randstarts[1] : randstarts[1] + randsize[1]] = randsquare
        init_mat = init_mat.to(torch.int16)

        return init_mat.to(self.device)  # (H,W)

torchgol_sub.py file to also be put in substrates/.
This is the file imitating the asal structure, but with pytorch. Here there are some clashes with torch OOP style, and JAX which I assume is much more functional.

"""
    An attempt of a simple wrapper of torchgol, with the convention of Asal substrates.
"""

from .ca2d import CA2D
import torch, torchvision as tv
from torchvision.transforms.transforms import Resize

class torchGOLSub():
    def __init__(self, grid_size, s_num:str="23", b_num:str="3", device='cpu'):
        """
            Not sure how devices are handled in jax, in pytorch we should pass it here, or later.
        """
        self.grid_size = grid_size
        self.device = device
        self.ca = CA2D(size=(grid_size, grid_size), b_num=b_num, s_num=s_num, device=device)

    def default_params(self, rng=None):
        """
            Not sure what the equivalent of rng should be in pytorch.
            Probably a torch.Generator, or the seed to generate the generator directly ?

            Also, here is unclear to me which format is supported for the parameters.
            Any dictionary with anything, I guess ?
        """

        # Here I generate the random rule directly here
        # In general, maybe it's best if it's a method of the underlying automaton class
        # directly, that way if there are constraints (e.g., we choose 3 channels, and so on)
        # they can be directly accounted for.
        b_rule = torch.randint(0, 2**9, (1,),generator=rng).item()
        s_rule = torch.randint(0, 2**9, (1,),generator=rng).item()

        return {'b_num': b_rule, 's_num': s_rule}

    def init_state(self, rng, params):
        """
            Here, would be nice to know in general what can 'state' be, is it completely free,
            or not ?
        """
        state = self.ca.get_init_mat(rand_portion=0.5) # (h, w) tensor
        return state # (H,W) tensor

    def step_state(self, rng, state, params):
        """
            For this function, I find it weird to have to pass 'params' here.
            My confusion comes probably because torch is much more OOP, while JAX seems 
            much more functional in nature. In my mind, we already have the initialize model
            (here self.ca, but e.g. in NCA, the initialize network), so does it make sense
             to have to re-initialize the model each time with params. Anyway, I keep it this way, but feels
            strange (e.g. in the NCA) to have to re-initialize the Neural network before each step, but might
            be standard in JAX ?
        """
        # Here I assume that the state is a 2D tensor, and that the parameters are a dictionary
        # with the keys 'b_num' and 's_num'.

        self.ca.change_num(params['s_num'], params['b_num'])
        # Here again, in 'non-functional' fashion, the ca state is stored in the ca object.
        # In principle I have no objection to having a method 'step_state' inside ca, and use that instead.
        self.ca.world = state 
        self.ca.step()

        
        return self.ca.world # (H,W) tensor

    def render_state(self, state, params, img_size=None):
        """
            Again for this function, the rendering takes only the state as input.
            This is a bit annoying because it prevents me from rendering things like
            'fading' of pixels, because then the rendering of the state depends also on
            the past of the state, not only the current state.^

            For now, I cheat and just render the current state of self.ca . 
            If this causes problems for some reason, one can use ca.pure_render_state() instead.
        """
        self.ca.draw() # renders the state
        img = self.ca._worldmap # (3, h, w) tensor
        # In functional fashion :
        # img = self.ca.pure_render_state(state) # (3, h, w) tensor
        
        if img_size is not None:
            img = Resize((img_size,img_size),interpolation=tv.transforms.InterpolationMode.NEAREST)(img)
        
        img = img.permute(1,2,0) # (3, h, w) -> (h, w, 3) to follow ASAL conventions

        return img # Note : the values are in [0.,1.] range

I've also added a small script that you can put in the root directory of asal/ just to check that everything works. It should generate a video of GoL in ./goltestoutput/videotensor.mp4.

import torch

if __name__ == '__main__':
    from showtens import save_video
    from substrates.torchgol_sub import torchGOLSub

    
    model = torchGOLSub(128, s_num="23", b_num="3", device='cuda')
    params = {'b_num': 8 , 's_num': 12} # GoL rule just to test
    state = model.init_state(rng=None, params=params)

    frames = []
    for i in range(100):
        state = model.step_state(None, state, params)
        img = model.render_state(state, params)
        frames.append(img)
    
    tensor_video = torch.stack(frames, dim=0) # (T, H, W, 3)
    tensor_video = tensor_video.permute(0,3,1,2) # (T, H, W, 3) -> (T, 3, H, W)

    save_video(tensor_video, 'goltestoutput', fps=15)

Let me know if any problem/questions arise with this example. If needed, I can also provide a Lenia implementation in pytorch, although it comes with a bit more overhead.

[fredericowieser]

Hey I have a version of the ASAL code in PyTorch that I made here, it is limited by only having the lenia substrate but should have all the other features there. It's not very well documented code but follows almost the exact same structure as this code so I hope it helps. https://github.com/fredericowieser/ASAL-PyTorch

The code comes from another project so has some artifacts from that, but I hope it still helps you anyway

[akarshkumar0101]

@fredericowieser Amazing work, thanks! I will make sure to add a link to that repo in our README!