rad-embeddings 0.1.0__tar.gz

rad_embeddings-0.1.0/.gitignore

@@ -0,0 +1,174 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# UV
+#   Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#uv.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/latest/usage/project/#working-with-version-control
+.pdm.toml
+.pdm-python
+.pdm-build/
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+# PyPI configuration file
+.pypirc
+
+# macos
+.DS_Store

rad_embeddings-0.1.0/.python-version

@@ -0,0 +1 @@
+3.12

rad_embeddings-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 RAD-Embeddings
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

rad_embeddings-0.1.0/PKG-INFO

@@ -0,0 +1,15 @@
+Metadata-Version: 2.4
+Name: rad-embeddings
+Version: 0.1.0
+Summary: Add your description here
+License-File: LICENSE
+Requires-Python: >=3.12
+Requires-Dist: dfa-gym>=0.1.0
+Requires-Dist: stable-baselines3>=2.4.1
+Requires-Dist: tensorboard>=2.18.0
+Requires-Dist: token-env>=0.2.1
+Requires-Dist: torch-geometric>=2.6.1
+Requires-Dist: torch<2.3.0
+Description-Content-Type: text/markdown
+
+# rad-embeddings

rad_embeddings-0.1.0/README.md

@@ -0,0 +1 @@
+# rad-embeddings

rad_embeddings-0.1.0/pyproject.toml

@@ -0,0 +1,18 @@
+[project]
+name = "rad-embeddings"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+requires-python = ">=3.12"
+dependencies = [
+    "dfa-gym>=0.1.0",
+    "stable-baselines3>=2.4.1",
+    "tensorboard>=2.18.0",
+    "token-env>=0.2.1",
+    "torch<2.3.0",
+    "torch-geometric>=2.6.1",
+]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"

rad_embeddings-0.1.0/rad_embeddings/__init__.py

@@ -0,0 +1 @@
+from encoder import Encoder

rad_embeddings-0.1.0/rad_embeddings/encoder.py

@@ -0,0 +1,97 @@
+import numpy as np
+from dfa import DFA
+from dfa_gym import DFAEnv
+from stable_baselines3 import PPO
+from utils import dfa2obs
+import gymnasium as gym
+from utils import DFAEnvFeaturesExtractor, LoggerCallback
+from stable_baselines3.common.callbacks import EvalCallback, StopTrainingOnNoModelImprovement
+
+class Encoder():
+    def __init__(self, load_file: str):
+        model = PPO.load(load_file)
+        model.set_parameters(load_file)
+        for param in model.policy.parameters():
+            param.requires_grad = False
+        model.policy.eval()
+        self.obs2rad = model.policy.features_extractor
+        self.rad2token = lambda _rad: model.policy.action_net(_rad).argmax(dim=1)
+        self.n_tokens = self.obs2rad.n_tokens
+
+    def dfa2rad(self, dfa: DFA) -> np.array:
+        assert len(dfa.inputs) == self.n_tokens
+        obs = dfa2obs(dfa)
+        rad = self.obs2rad(obs)
+        return rad
+
+    @staticmethod
+    def train(n_tokens: int, train_env: gym.Env, save_dir: str, eval_env: gym.Env | None = None, id: str = "rad"):
+        save_dir = save_dir[:-1] if save_dir.endswith("/") else save_dir
+        config = dict(
+            policy = "MlpPolicy",
+            env = train_env,
+            learning_rate = 1e-3,
+            n_steps = 512,
+            batch_size = 1024,
+            n_epochs = 2,
+            gamma = 0.9,
+            gae_lambda = 0.5,
+            clip_range = 0.1,
+            ent_coef = 1e-2,
+            vf_coef = 0.5,
+            max_grad_norm = 0.5,
+            policy_kwargs = dict(
+                features_extractor_class = DFAEnvFeaturesExtractor,
+                features_extractor_kwargs = dict(features_dim=32, n_tokens=n_tokens),
+                net_arch=dict(pi=[], vf=[]),
+                share_features_extractor=True,
+            ),
+            verbose = 10,
+            tensorboard_log = f"{save_dir}/runs/"
+        )
+
+        model = PPO(**config)
+
+        print("Total number of parameters:", sum(p.numel() for p in model.policy.parameters() if p.requires_grad))
+        print(model.policy)
+
+        callback_list = []
+        logger_callback = LoggerCallback(gamma=config["gamma"])
+        callback_list.append(logger_callback)
+        if eval_env is not None:
+            stop_train_callback = StopTrainingOnNoModelImprovement(max_no_improvement_evals=10, min_evals=20, verbose=1)
+            eval_callback = EvalCallback(eval_env, eval_freq=1000, callback_after_eval=stop_train_callback, verbose=1)
+            callback_list.append(eval_callback)
+        model.learn(10_000_000, callback=callback_list)
+        model.save(f"{save_dir}/{id}")
+
+if __name__ == "__main__":
+    from stable_baselines3.common.env_util import make_vec_env
+    from stable_baselines3.common.env_checker import check_env
+
+    n_envs = 16
+    env_id = "DFAEnv-v1"
+    encoder_id = env_id + "-encoder"
+    save_dir = "storage"
+
+    env = gym.make(env_id)
+    check_env(env)
+
+    n_tokens = env.unwrapped.sampler.n_tokens
+
+    train_env = make_vec_env(env_id, n_envs=n_envs)
+    eval_env = gym.make(env_id)
+
+    Encoder.train(n_tokens=n_tokens, train_env=train_env, eval_env=eval_env, save_dir=save_dir, id=encoder_id)
+
+    sampler = env.unwrapped.sampler
+    encoder = Encoder(load_file=f"{save_dir}/{encoder_id}")
+
+    dfa = sampler.sample()
+    print(dfa)
+
+    rad = encoder.dfa2rad(dfa)
+    print(rad)
+
+    token = encoder.rad2token(rad)
+    print(token)

rad_embeddings-0.1.0/rad_embeddings/model.py

@@ -0,0 +1,28 @@
+import torch
+import torch.nn as nn
+from torch_geometric.nn import GATv2Conv
+
+class Model(nn.Module):
+    def __init__(self, input_dim, output_dim, **kwargs):
+        super().__init__()
+        self.input_dim = input_dim
+        self.output_dim = output_dim
+        self.hidden_dim = kwargs.get("hidden_dim", 64)
+        self.num_layers = kwargs.get("num_layers", 8)
+        self.n_heads = kwargs.get("n_heads", 4)
+        self.linear_in = nn.Linear(self.input_dim, self.hidden_dim)
+        self.conv = GATv2Conv(2*self.hidden_dim, self.hidden_dim, heads=self.n_heads, add_self_loops=False)
+        self.activation = nn.Tanh()
+        self.g_embed = nn.Linear(self.hidden_dim, self.output_dim)
+
+    def forward(self, data):
+        feat = data.feat
+        edge_index = data.edge_index
+        current_state = data.current_state
+        h_0 = self.linear_in(feat.float())
+        h = h_0
+        for i in range(self.num_layers):
+            h = self.conv(torch.cat([h, h_0], dim=1), edge_index).view(h.shape[0], self.n_heads, self.hidden_dim).sum(dim=1)
+            h = self.activation(h)
+        hg = h[current_state.bool()]
+        return self.g_embed(hg)

rad_embeddings-0.1.0/rad_embeddings/train_token_env_policy.py

@@ -0,0 +1,64 @@
+import torch
+import token_env
+import gymnasium as gym
+from encoder import Encoder
+from dfa_gym import DFAWrapper
+from stable_baselines3 import PPO
+from utils import TokenEnvFeaturesExtractor, LoggerCallback
+from stable_baselines3.common.env_util import make_vec_env
+from stable_baselines3.common.env_checker import check_env
+from dfa_samplers import ReachSampler, ReachAvoidSampler, RADSampler
+
+# import wandb
+# from wandb.integration.sb3 import WandbCallback
+
+# run = wandb.init(project="sb3", sync_tensorboard=True)
+
+n_envs = 16
+env_id = "TokenEnv-v0"
+
+env = gym.make(env_id)
+check_env(env)
+n_tokens = env.unwrapped.n_tokens
+
+# env_kwargs = dict(env_id = "TokenEnv-v0", sampler=RADSampler(n_tokens=n_tokens), label_f=token_env.TokenEnv.label_f)
+env_kwargs = dict(env_id = "TokenEnv-v0", sampler=ReachAvoidSampler(n_tokens=n_tokens, max_size=4, prob_stutter=1.0), label_f=token_env.TokenEnv.label_f)
+env = make_vec_env(DFAWrapper, env_kwargs=env_kwargs, n_envs=n_envs)
+
+encoder = Encoder(load_file="storage/DFAEnv-v0-encoder")
+
+config = dict(
+    policy = "MultiInputPolicy",
+    env = env,
+    n_steps = 128,
+    batch_size = 256,
+    gamma = 0.99,
+    policy_kwargs = dict(
+        features_extractor_class=TokenEnvFeaturesExtractor,
+        features_extractor_kwargs=dict(features_dim=1056, encoder=encoder),
+        net_arch=dict(pi=[64, 64, 64], vf=[64, 64]),
+        share_features_extractor=True,
+        activation_fn=torch.nn.ReLU
+    ),
+    verbose = 10,
+    tensorboard_log = f"token_env_reach_avoid_policy/runs/"
+    # tensorboard_log = f"token_env_reach_avoid_policy/runs/{run.id}"
+)
+
+model = PPO(**config)
+
+print("Total number of parameters:", sum(p.numel() for p in model.policy.parameters() if p.requires_grad))
+print(model.policy)
+
+logger_callback = LoggerCallback(gamma=config["gamma"])
+# wandb_callback = WandbCallback(
+#     gradient_save_freq=100,
+#     model_save_freq=100,
+#     model_save_path=f"token_env_reach_avoid_policy/models/{run.id}",
+#     verbose=config["verbose"])
+
+# model.learn(10_000_000, callback=[logger_callback, wandb_callback])
+model.learn(1_000_000, callback=[logger_callback])
+model.save("token_env_reach_avoid_policy/token_env_reach_avoid_policy")
+
+# run.finish()

rad_embeddings-0.1.0/rad_embeddings/utils/__init__.py

@@ -0,0 +1,4 @@
+from utils.sb3_dfa_env_features_extractor import DFAEnvFeaturesExtractor
+from utils.sb3_token_env_features_extractor import TokenEnvFeaturesExtractor
+from utils.sb3_logger_callback import LoggerCallback
+from utils.utils import obs2feat, dfa2obs, dfa2dist

rad_embeddings-0.1.0/rad_embeddings/utils/sb3_dfa_env_features_extractor.py

@@ -0,0 +1,14 @@
+from model import Model
+from utils.utils import feature_inds, obs2feat
+from stable_baselines3.common.torch_layers import BaseFeaturesExtractor
+
+class DFAEnvFeaturesExtractor(BaseFeaturesExtractor):
+    def __init__(self, observation_space, features_dim, model_cls=Model, n_tokens=10):
+        super().__init__(observation_space, features_dim)
+        in_feat_size = n_tokens + len(feature_inds)
+        self.model = model_cls(in_feat_size, features_dim)
+        self.n_tokens = n_tokens
+
+    def forward(self, obs):
+        feat = obs2feat(obs, n_tokens=self.n_tokens)
+        return self.model(feat)

rad_embeddings-0.1.0/rad_embeddings/utils/sb3_logger_callback.py

@@ -0,0 +1,50 @@
+import numpy as np
+from collections import deque
+from stable_baselines3.common.callbacks import BaseCallback
+
+class LoggerCallback(BaseCallback):
+    def __init__(self, gamma=0.9):
+        super().__init__()
+        self.gamma = gamma
+        self.n = 100
+        self.rewards = deque(maxlen=self.n)
+        self.discounted_rewards = deque(maxlen=self.n)
+        self.episode_lengths = deque(maxlen=self.n)
+
+    def _on_step(self):
+        # Log scalar value (here a random variable)
+        dones = self.locals["dones"]
+        infos = self.locals["infos"]
+
+        for idx, done in enumerate(dones):
+            info = infos[idx]
+            if (
+                done
+                and info.get("episode") is not None
+            ):
+                episode_info = info["episode"]
+                reward = info["episode"]['r']
+                episode_length = info["episode"]['l']
+                t = episode_length - 1
+                discounted_reward = reward * (self.gamma ** t)
+                self.rewards.append(reward)
+                self.discounted_rewards.append(discounted_reward)
+                self.episode_lengths.append(episode_length)
+
+        return True
+
+    def _on_rollout_end(self):
+        ep_len_min = np.min(self.episode_lengths)
+        ep_len_max = np.max(self.episode_lengths)
+        ep_len_std = np.std(self.episode_lengths)
+        ep_rew_min = np.min(self.rewards)
+        ep_rew_max = np.max(self.rewards)
+        ep_rew_std = np.std(self.rewards)
+        ep_disc_rew_mean = np.mean(self.discounted_rewards)
+        self.logger.record("rollout/ep_len_min", ep_len_min)
+        self.logger.record("rollout/ep_len_max", ep_len_max)
+        self.logger.record("rollout/ep_len_std", ep_len_std)
+        self.logger.record("rollout/ep_rew_min", ep_rew_min)
+        self.logger.record("rollout/ep_rew_max", ep_rew_max)
+        self.logger.record("rollout/ep_rew_std", ep_rew_std)
+        self.logger.record("rollout/ep_rew_disc_mean", ep_disc_rew_mean)

rad_embeddings-0.1.0/rad_embeddings/utils/sb3_token_env_features_extractor.py

@@ -0,0 +1,27 @@
+import torch
+from torch import nn
+from model import Model
+from stable_baselines3.common.torch_layers import BaseFeaturesExtractor
+
+class TokenEnvFeaturesExtractor(BaseFeaturesExtractor):
+    def __init__(self, observation_space, features_dim, encoder):
+        super().__init__(observation_space, features_dim)
+        self.encoder = encoder
+        c, w, h = observation_space["obs"].shape # CxWxH
+        self.image_conv = nn.Sequential(
+            nn.Conv2d(c, 16, (2, 2)),
+            nn.ReLU(),
+            nn.Conv2d(16, 32, (2, 2)),
+            nn.ReLU(),
+            nn.Conv2d(32, 64, (2, 2)),
+            nn.ReLU(),
+            nn.Flatten()
+        )
+
+    def forward(self, dict_obs):
+        dfa_obs = dict_obs["dfa_obs"]
+        obs = dict_obs["obs"]
+        rad = self.encoder.obs2rad(dfa_obs)
+        obs = self.image_conv(obs)
+        obs = torch.cat((obs, rad), dim=1)
+        return obs

rad_embeddings-0.1.0/rad_embeddings/utils/utils.py

@@ -0,0 +1,85 @@
+import torch
+import numpy as np
+from dfa import DFA
+from dfa.utils import dfa2dict
+from stable_baselines3 import PPO
+from collections import OrderedDict
+from torch_geometric.data import Data
+from torch_geometric.data import Batch
+from dfa.utils import min_distance_to_accept_by_state
+
+feature_inds = {"temp": -5, "rejecting": -4, "accepting": -3, "init": -2, "normal": -1}
+
+def obs2feat(dfa_obs, n_tokens):
+    if dfa_obs.ndim == 1:
+        return _obs2feat(dfa_obs, n_tokens=n_tokens)
+    elif dfa_obs.ndim == 2:
+        return Batch.from_data_list(list(map(lambda x: _obs2feat(x, n_tokens=n_tokens), dfa_obs)))
+    else:
+        raise ValueError(f"Invalid ndim for dfa_obs: expected 1 or 2, but got {dfa_obs.ndim}")
+
+def _obs2feat(dfa_obs, n_tokens):
+    tokens = list(range(n_tokens))
+    feature_size = len(tokens) + len(feature_inds)
+    dfa_int = int("".join(map(str, map(int, dfa_obs.squeeze().tolist()))))
+    dfa = DFA.from_int(dfa_int, tokens)
+    dfa_dict, s_init = dfa2dict(dfa)
+    nodes = OrderedDict({s: np.zeros(feature_size) for s in dfa_dict.keys()})
+    if len(nodes) == 1:
+        edges = [(0, 0)]
+    else:
+        edges = [(s, s) for s in nodes]
+    for s in dfa_dict.keys():
+        label, transitions = dfa_dict[s]
+        leaving_transitions = [1 if s != transitions[a] else 0 for a in transitions.keys()]
+        if s not in nodes:
+            nodes[s] = np.zeros(feature_size)
+        nodes[s][feature_inds["normal"]] = 1
+        if s == s_init:
+            nodes[s][feature_inds["init"]] = 1
+        if label: # is accepting?
+            nodes[s][feature_inds["accepting"]] = 1
+        elif sum(leaving_transitions) == 0: # is rejecting?
+            nodes[s][feature_inds["rejecting"]] = 1
+        for e in dfa_dict.keys():
+            if s == e:
+                continue
+            for a in transitions:
+                if transitions[a] == e:
+                    if (s, e) not in nodes:
+                        nodes[(s, e)] = np.zeros(feature_size)
+                        nodes[(s, e)][feature_inds["temp"]] = 1
+                    nodes[(s, e)][a] = 1
+                    s_idx = list(nodes.keys()).index(s)
+                    t_idx = list(nodes.keys()).index((s, e))
+                    e_idx = list(nodes.keys()).index(e)
+                    # Reverse
+                    if (e_idx, t_idx) not in edges:
+                        edges.append((e_idx, t_idx))
+                    if (t_idx, t_idx) not in edges:
+                        edges.append((t_idx, t_idx))
+                    if (t_idx, s_idx) not in edges:
+                        edges.append((t_idx, s_idx))
+    feat = torch.from_numpy(np.array(list(nodes.values())))
+    edge_index = torch.from_numpy(np.array(edges))
+    current_state = torch.from_numpy(np.array([1] + [0] * (len(nodes) - 1))) # 0 is the current state
+    return Data(feat=feat, edge_index=edge_index.T, current_state=current_state)
+
+def dfa2obs(dfa: DFA) -> Data:
+    return np.array([int(i) for i in str(dfa.to_int())])
+
+def dfa2dist(dfa_obs, n_tokens):
+    tokens = list(range(n_tokens))
+    feature_size = len(tokens) + len(feature_inds)
+
+    dfa_int = int("".join(map(str, map(int, dfa_obs.squeeze().tolist()))))
+    dfa = DFA.from_int(dfa_int, tokens)
+
+    n = len(dfa.states())
+
+    dist = min_distance_to_accept_by_state(dfa)[0]
+
+    if dist == float("inf"):
+        dist = 0
+
+    return dist