GReinSS: Generative RL for Structured States

• GReinSS is a framework for learning distributions over discrete, mechanistic latent structures such as graphs, subsets, and RNA isoforms from indirect observations.
• It employs a policy-gradient approach with dynamic reward rescaling, operating directly in the combinatorial latent space without resorting to continuous latent codes.
• Empirical results demonstrate its superior performance over baselines across simulated graphs, subsets, and RNA isoform reconstruction, highlighting its robustness and wide applicability.

GReinSS, short for Generative Reinforcement Learning of Structured States, is a framework for learning a distribution over discrete, mechanistic latent structures such as sets, graphs, and RNA isoforms from indirect observations generated through a known or modeled likelihood $\Pr(X \mid S)$. The framework addresses a latent-variable maximum-likelihood problem in which the latent state $S$ is never observed directly, but observations $X$ are available, and the objective is to learn a generative model $\Pr_\theta(S)$ that matches the unknown true latent-state distribution by maximizing the observed-data likelihood (Ivanovic et al., 5 Jun 2026). Its central methodological contribution is a dynamic reward rescaling within a policy-gradient formulation, yielding an unbiased estimator of $\nabla_\theta \log \Pr_\theta(X_{1:N})$ while operating directly in the combinatorial latent space rather than in an artificial continuous code space (Ivanovic et al., 5 Jun 2026).

1. Problem formulation and conceptual scope

GReinSS is designed for settings in which latent structures $S \in \mathcal{S}$ are drawn from an unknown true distribution $\Pr^*(S)$, but only indirect observations $X$ are available, with $X \sim \Pr(X \mid S)$. The task is to learn a parametric generative model

$\Pr_\theta(S) = \Pr(S \mid \theta)$

such that the marginal likelihood

$S$0

is maximized over the observed dataset (Ivanovic et al., 5 Jun 2026).

The framework targets combinatorial latent spaces in which $S$1 may consist of all subsets of a large universe, all graphs over many nodes, or all isoforms over a gene. In such settings, classical expectation maximization is described as struggling because the E-step requires expectations over an exponentially large state space unless a special structure such as dynamic programming is available (Ivanovic et al., 5 Jun 2026). Deep latent-variable methods such as variational autoencoders are contrasted with this approach because they generally introduce artificial continuous latent codes in $S$2 rather than reconstructing the mechanistic structures themselves (Ivanovic et al., 5 Jun 2026).

A defining feature of GReinSS is that the latent space is the actual combinatorial structure space of interest. The method therefore combines generative modeling, latent-state inference, and structured policy learning in a single objective. This suggests that its notion of representation is tied to domain objects rather than to a purely statistical embedding.

2. Likelihood maximization through dynamic policy gradients

For a dataset $S$3, GReinSS assumes conditional independence given $S$4, so that

$S$5

If $S$6 denotes a trajectory that constructs a latent state $S$7, then the marginal for one observation can be written as

$S$8

where $S$9 (Ivanovic et al., 5 Jun 2026).

Using the log-derivative trick, the paper derives

$X$0

This yields the dynamic reward

$X$1

with the resulting identity

$X$2

The reward depends on $X$3 through the denominators $X$4, but the method does not backprop through the reward; $X$5 is treated as a scalar in the gradient step, and the estimator remains unbiased for the likelihood objective (Ivanovic et al., 5 Jun 2026).

The conceptual difference from a naive reward is central. A naive construction,

$X$6

favors trajectories with high raw explanatory probability and can collapse onto a small number of easy-to-explain states. By contrast, the dynamic denominator down-weights observations already well explained by the current model and up-weights those with small current marginal likelihood. In the formulation of the paper, this causes the policy to optimize the product $X$7 rather than overcommitting to a single high-reward trajectory (Ivanovic et al., 5 Jun 2026). A plausible implication is that the reward acts as an adaptive balancing term across heterogeneous observations.

Mini-batching is accommodated by defining, for a batch $X$8,

$X$9

and showing that

$\Pr_\theta(S)$0

so that a suitably scaled mini-batch reward still yields an unbiased estimator of the full gradient (Ivanovic et al., 5 Jun 2026).

3. Policy representation, off-policy proposals, and optimization procedure

GReinSS represents the generative model $\Pr_\theta(S)$1 as a policy over trajectories. A trajectory

$\Pr_\theta(S)$2

is a sequence of discrete actions that incrementally builds a terminal state $\Pr_\theta(S)$3. The policy factorizes as

$\Pr_\theta(S)$4

where $\Pr_\theta(S)$5 is implemented by a neural network (Ivanovic et al., 5 Jun 2026). The policy network used in the reported experiments is always a small 2-layer MLP (50 hidden units, leakyReLU), although the meaning of actions and partial states depends on the application (Ivanovic et al., 5 Jun 2026).

The training loop consists of initializing $\Pr_\theta(S)$6, repeatedly sampling trajectories, estimating marginals $\Pr_\theta(S)$7 by Monte Carlo, computing rewards, and applying a policy-gradient update until convergence of $\Pr_\theta(S)$8 (Ivanovic et al., 5 Jun 2026). With on-policy sampling,

$\Pr_\theta(S)$9

and with off-policy sampling from $\nabla_\theta \log \Pr_\theta(X_{1:N})$0,

$\nabla_\theta \log \Pr_\theta(X_{1:N})$1

followed by gradient ascent $\nabla_\theta \log \Pr_\theta(X_{1:N})$2 (Ivanovic et al., 5 Jun 2026). Reward normalization, baselines, gradient clipping, entropy regularization, Adam or SGD, and choices of batch size and off-policy bias strength are treated as auxiliary implementation mechanisms rather than as the core innovation (Ivanovic et al., 5 Jun 2026).

The paper also derives the variance-minimizing off-policy proposal

$\nabla_\theta \log \Pr_\theta(X_{1:N})$3

where

$\nabla_\theta \log \Pr_\theta(X_{1:N})$4

Equivalently,

$\nabla_\theta \log \Pr_\theta(X_{1:N})$5

This proposal is a mixture of posteriors over trajectories induced by individual observations (Ivanovic et al., 5 Jun 2026). Exact sampling from this $\nabla_\theta \log \Pr_\theta(X_{1:N})$6 is described as difficult, so the implementation uses heuristic biased proposals that move sampling toward that target, with importance-sampling correction. This suggests a close relationship between GReinSS and posterior-focused Monte Carlo estimation, even though the framework is presented as a policy-learning method.

4. Instantiations for graphs, subsets, and RNA isoforms

GReinSS is structurally agnostic provided that a latent object can be constructed as a sequence of discrete actions (Ivanovic et al., 5 Jun 2026). The paper details three instantiations.

Latent directed graphs

For process graphs, each latent state is a directed graph over a fixed node set, represented as a set of directed edges. The policy starts from the empty graph and, at each step, selects either a directed edge to add or a termination action (Ivanovic et al., 5 Jun 2026). Observations consist of lists of start and end nodes of absorbing random walks on the graph. For a graph $\nabla_\theta \log \Pr_\theta(X_{1:N})$7, the probability of a random walk from node $\nabla_\theta \log \Pr_\theta(X_{1:N})$8 to node $\nabla_\theta \log \Pr_\theta(X_{1:N})$9 is computed via the inverse shifted Laplacian $S \in \mathcal{S}$0, following Wu et al. 2012, and $S \in \mathcal{S}$1 is the product over the random walks in $S \in \mathcal{S}$2 of these start-end probabilities (Ivanovic et al., 5 Jun 2026).

Latent subsets of a universe

For subset inference, the latent state is a subset $S \in \mathcal{S}$3, equivalently a binary vector in $S \in \mathcal{S}$4. The policy starts with the empty set and sequentially adds elements, together with a termination action (Ivanovic et al., 5 Jun 2026). Observations are noisy real vectors $S \in \mathcal{S}$5, with coordinate-wise likelihood

$S \in \mathcal{S}$6

Hence

$S \in \mathcal{S}$7

For this task the implementation also includes off-policy sampling tailored to the Gaussian likelihood, biasing addition probabilities according to a log-probability ratio comparing $S \in \mathcal{S}$8 and $S \in \mathcal{S}$9 (Ivanovic et al., 5 Jun 2026).

RNA isoforms from short-read RNA-seq

For isoform reconstruction, latent structures are RNA isoforms for each gene and sequencing sample. Isoforms are represented as sequences of junctions, where a junction is a pair specifying the end of one exon and the start of the next (Ivanovic et al., 5 Jun 2026). The set of possible exons and known junctions is taken from GENCODE v39. A latent isoform state contains a chain of junctions defining a transcript, the sample index, and the genomic position of the read when a particular read is embedded in the state representation (Ivanovic et al., 5 Jun 2026).

During training, the policy starts with an empty isoform, sequentially selects legal next junctions subject to exon connectivity constraints, chooses a sequencing sample via a small network, and uses a learned bias vector to specify the probability of emitting reads from each junction of the isoform (Ivanovic et al., 5 Jun 2026). The input data are short-read RNA-seq junction-overlapping read counts from STAR. Each read observation includes a sequencing-sample index and a genomic position indicating the junction it overlaps. Under the state construction used in the model,

$\Pr^*(S)$0

if the isoform/sample/junction in the state matches the read and

$\Pr^*(S)$1

otherwise (Ivanovic et al., 5 Jun 2026). In consequence, the likelihood acts as a mask over state-read compatibility, while isoform proportions and junction-level read probabilities are modeled through the policy.

5. Relation to EM, VAEs, GFlowNets, and other baselines

The paper positions GReinSS against three families of alternatives: classical or approximate EM, deep latent-variable models such as VAEs, and policy-learning baselines (Ivanovic et al., 5 Jun 2026).

Classical EM is characterized as requiring either per-data posteriors $\Pr^*(S)$2 or a distribution

$\Pr^*(S)$3

in the E-step, followed by an M-step maximizing the expected complete-data log-likelihood (Ivanovic et al., 5 Jun 2026). In combinatorial spaces this is intractable in general. The implemented generalized EM (GEM) baselines therefore use hard assignments

$\Pr^*(S)$4

and then maximize

$\Pr^*(S)$5

using VAE, autoregressive, or discrete diffusion models in the M-step (Ivanovic et al., 5 Jun 2026). These baselines are described as limited by their reliance on hard assignments and their failure to properly marginalize over $\Pr^*(S)$6, especially when the posterior is multimodal or uncertain (Ivanovic et al., 5 Jun 2026).

VAEs are contrasted with GReinSS on three dimensions: latent space, inference mechanism, and training objective. In VAEs, one typically introduces a continuous code $\Pr^*(S)$7, learns an encoder $\Pr^*(S)$8, and optimizes an ELBO surrogate. In GReinSS, by contrast, the latent states are the mechanistic objects themselves, no variational posterior $\Pr^*(S)$9 is used, and the objective is the direct maximization of

$X$0

under the specified observation model $X$1 (Ivanovic et al., 5 Jun 2026). The paper attributes to this design improved interpretability and the possibility of improved identifiability when the observation model encodes mechanistic knowledge (Ivanovic et al., 5 Jun 2026).

Among policy-learning baselines, the paper compares against naive policy gradients with reward $X$2 and GFlowNets trained with trajectory balance using the same reward proxy (Ivanovic et al., 5 Jun 2026). Naive policy gradients are reported to collapse to degenerate solutions, such as predicting the empty graph, while GFlowNets are noted to be optimal only in a special case in which each observation is explained by exactly one trajectory (Ivanovic et al., 5 Jun 2026). In the general combinatorial mixture setting, they optimize a proxy rather than the marginal likelihood. Local search is also included as a non-RL baseline, performing hill-climbing on $X$3 for each observation while ignoring $X$4 (Ivanovic et al., 5 Jun 2026).

Theoretical special cases further clarify these relationships. If observations equal states and $X$5 is one only when $X$6, the problem reduces to standard maximum-likelihood generative modeling and GReinSS collapses to standard autoregressive training. If each observation has exactly one feasible latent state, local search suffices. If $X$7 is independent of $X$8, the denominator cancels and GReinSS reduces to standard policy gradients. Under a particular one-trajectory-per-observation setup, GFlowNets are optimal for the problem formulation (Ivanovic et al., 5 Jun 2026).

6. Empirical results, theoretical implications, and limitations

The empirical study covers simulated latent graphs, simulated latent sets, and RNA isoform reconstruction from real short-read RNA-seq data (Ivanovic et al., 5 Jun 2026).

For simulated latent graphs, the latent graph is generated by thresholding an Erdős–Rényi base directed graph with $X$9 and edge weights sampled from $X \sim \Pr(X \mid S)$0; each observation consists of $X \sim \Pr(X \mid S)$1 absorbing random walks with $X \sim \Pr(X \mid S)$2 (Ivanovic et al., 5 Jun 2026). The evaluation metric is F1 score on edge-set reconstruction. Across all $X \sim \Pr(X \mid S)$3, GReinSS achieves the highest median F1; for the low-information setting $X \sim \Pr(X \mid S)$4, the reported median F1 is approximately $X \sim \Pr(X \mid S)$5 versus less than $X \sim \Pr(X \mid S)$6 for all baselines (Ivanovic et al., 5 Jun 2026). Naive policy gradients often predict empty graphs.

For simulated latent sets, the universe size is varied over $X \sim \Pr(X \mid S)$7 and the noise level over $X \sim \Pr(X \mid S)$8 (Ivanovic et al., 5 Jun 2026). Latent sets are generated from reusable modules in a dictionary-learning-like construction, and F1 on set membership is used for evaluation. For $X \sim \Pr(X \mid S)$9, GReinSS has the highest median F1 at all noise levels. At low noise, local search and naive policy or GFlowNets are competitive because each observation almost uniquely identifies the latent set. At higher noise, GEM-based methods become the strongest baselines, but remain below GReinSS (Ivanovic et al., 5 Jun 2026). With $\Pr_\theta(S) = \Pr(S \mid \theta)$0 and $\Pr_\theta(S) = \Pr(S \mid \theta)$1, the reported median F1 values are approximately $\Pr_\theta(S) = \Pr(S \mid \theta)$2 for GReinSS, $\Pr_\theta(S) = \Pr(S \mid \theta)$3 for local search, $\Pr_\theta(S) = \Pr(S \mid \theta)$4 for naive policy gradients, and less than $\Pr_\theta(S) = \Pr(S \mid \theta)$5 for GEM-based VAE, autoregressive, and diffusion models (Ivanovic et al., 5 Jun 2026). Off-policy sampling is reported as important for all policy-based methods, while GReinSS is relatively robust to a $\Pr_\theta(S) = \Pr(S \mid \theta)$6 change in off-policy bias strength (Ivanovic et al., 5 Jun 2026).

For RNA isoform reconstruction, the evaluation uses GTEx short-read junction-overlapping counts across multiple human tissue samples, with long-read RNA-seq available for 61 of 17,371 samples to provide isoforms and sample-specific proportions via FLAIR (Ivanovic et al., 5 Jun 2026). Predicted and long-read isoforms are represented as covered genomic regions between first and last junction, and comparison is performed using a Jaccard-distance cost matrix within an optimal-transport problem. Sample-wise transport costs are then averaged over the 61 samples, weighted by long-read isoform read counts (Ivanovic et al., 5 Jun 2026).

The reported MBD2 example illustrates the magnitude of the difference. In one fibroblast sample, long-read FLAIR identifies two isoforms with proportions 96% and 4%, and a Jaccard distance of 0.1495 between them. GReinSS predicts two matching isoforms with proportions 92% and 7%, giving an error of approximately 0.0067 for that sample, whereas RSEM predicts the major isoform at only 18% and spreads the remaining mass across very different isoforms, giving an error of approximately 0.5374 (Ivanovic et al., 5 Jun 2026). Aggregated over 61 samples with weighting by long-read counts, the reported errors are approximately 0.0069 for GReinSS and 0.471 for RSEM, with a difference of approximately $\Pr_\theta(S) = \Pr(S \mid \theta)$7 (Ivanovic et al., 5 Jun 2026). Across 14,390 genes, GReinSS beats RSEM by at least 0.05 on 46.6% of genes, while RSEM beats GReinSS by at least 0.05 on only 9.4% (Ivanovic et al., 5 Jun 2026). On a subset of genes comparing additional policy baselines, the median errors are reported as 0.141 for GReinSS, 0.317 for RSEM, 0.388 for GFlowNets, and 0.427 for naive policy gradients (Ivanovic et al., 5 Jun 2026).

Theoretically, the paper emphasizes three points: the exact likelihood-gradient identity induced by the dynamic reward, the optimal off-policy proposal as a mixture of posteriors, and the reduction to classical methods in identifiable special cases (Ivanovic et al., 5 Jun 2026). At the same time, it notes several limitations: scalability due to repeated trajectory sampling and evaluation of $\Pr_\theta(S) = \Pr(S \mid \theta)$8; the assumption that a forward observation model is known or easily approximated; the high variance characteristic of REINFORCE-style gradients despite dynamic rewards and off-policy sampling; and the fact that only relatively simple fully connected MLP policies are explored, leaving richer GNN or transformer-based policies for more complex structure spaces unexamined (Ivanovic et al., 5 Jun 2026).

Within a broader research context, GReinSS lies at the intersection of RL for generative modeling, discrete latent-variable models, and structured prediction with search (Ivanovic et al., 5 Jun 2026). Potential future applications explicitly mentioned or implied include regulatory networks, phylogenetic trees, chemical reaction pathways, transportation or communication networks, and program induction or synthesis; methodological extensions include actor-critic or Q-learning variants, auxiliary networks for approximating the optimal off-policy proposal, joint learning of parts of $\Pr_\theta(S) = \Pr(S \mid \theta)$9, and richer multi-environment or hierarchical latent distributions (Ivanovic et al., 5 Jun 2026). This suggests that the framework is less a task-specific algorithm than a general recipe for likelihood-based learning over combinatorial latent state spaces.