This repo contains a JAX implementation of RAD embeddings, see project webpage for more information.

Installation

Install using pip.

pip install rad-embeddings # CPU-only
pip install rad-embeddings[cuda] # With CUDA

Usage

Instantiate a pretrained encoder.

from rad_embeddings import Encoder

encoder = Encoder() # Loads a pretrained DFA encoder with default parameters: handles at most 10-state DFAs with 10-token alphabets

Use DFAx to sample DFAs.

import jax
from dfax.samplers import ReachSampler, ReachAvoidSampler, RADSampler

sampler = RADSampler()
key = jax.random.PRNGKey(42)

key, subkey = jax.random.split(key)
dfax = sampler.sample(subkey)

Pass the DFA to the encoder to get its embedding — both DFAx and DFA objects are supported.

from dfax import dfax2dfa

rad_embed_from_dfax = encoder(dfax)
rad_embed_from_dfa = encoder(dfax2dfa(dfax)) # Returns the same embedding

Compute bisimulation distance between two DFA embeddings.

key, subkey = jax.random.split(key)
dfax_l = sampler.sample(subkey)
rad_l = encoder(dfax_l)

key, subkey = jax.random.split(key)
dfax_r = sampler.sample(subkey)
rad_r = encoder(dfax_r)

distance = encoder.distance(rad_l, rad_r)

Solve a one-step bisimulation problem.

from dfax import DFAx
import jax.numpy as jnp

# Reach token 1 and then token 2 while avoding token 9
dfa_l = DFAx.create(
	start = 0,
	transitions = jnp.array([
		[0, 1, 0, 0, 0, 0, 0, 0, 0, 3],
		[1, 1, 2, 1, 1, 1, 1, 1, 1, 3],
		[2, 2, 2, 2, 2, 2, 2, 2, 2, 2],
		[3, 3, 3, 3, 3, 3, 3, 3, 3, 3],
		[4, 4, 4, 4, 4, 4, 4, 4, 4, 4],
	]),
	labels = jnp.array([False, False, True, False, False])
)

# Reach token 9
dfa_r = DFAx.create(
	start = 0,
	transitions = jnp.array([
		[0, 0, 0, 0, 0, 0, 0, 0, 0, 1],
		[1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
		[2, 2, 2, 2, 2, 2, 2, 2, 2, 2],
		[3, 3, 3, 3, 3, 3, 3, 3, 3, 3],
		[4, 4, 4, 4, 4, 4, 4, 4, 4, 4],
	]),
	labels = jnp.array([False, True, False, False, False])
)

distinguishing_action = encoder.solve(dfa_l, dfa_r) # Returns 1 as token 1 is the one-step distinguishing action

Train a new encoder — see encoder.py for the default arguments of the EncoderModule.train class method.

from rad_embeddings import EncoderModule

EncoderModule.train(max_size=5, n_tokens=5, debug=True, save_dir="my_storage")

Load your trained parameters.

from rad_embeddings import Encoder

encoder = Encoder(max_size=5, n_tokens=5, storage_dir="my_storage")

Citation

Please cite the following papers if you use RAD Embeddings in your work.

@inproceedings{DBLP:conf/nips/YalcinkayaLVS24,
  author       = {Beyazit Yalcinkaya and
                  Niklas Lauffer and
                  Marcell Vazquez{-}Chanlatte and
                  Sanjit A. Seshia},
  title        = {Compositional Automata Embeddings for Goal-Conditioned Reinforcement
                  Learning},
  booktitle    = {NeurIPS},
  year         = {2024}
}

@inproceedings{DBLP:conf/neus/YalcinkayaLVS25,
  author       = {Beyazit Yalcinkaya and
                  Niklas Lauffer and
                  Marcell Vazquez{-}Chanlatte and
                  Sanjit A. Seshia},
  title        = {Provably Correct Automata Embeddings for Optimal Automata-Conditioned
                  Reinforcement Learning},
  booktitle    = {NeuS},
  series       = {Proceedings of Machine Learning Research},
  volume       = {288},
  pages        = {661--675},
  publisher    = {{PMLR}},
  year         = {2025}
}