1. Introduction
Artificial intelligence (AI) is reshaping many industries and is changing the way that we use tools to aid us in our everyday lives. Engineering design is no different, with new AI tools being used in many areas ranging from the conceptual design of buildings (Reference Roy, Kim, Yin, Yeh, Nakabayashi, Campbell, Keough and TsujiRoy et al., 2022) to AI options in computer-aided design (CAD) tools (AutoDesk, 2025). In particular, the recent surge in the use of generative AI has opened interesting avenues for generative modelling within the domain of design.
A key challenge of engineering design is data representation and generation. This challenge distinguishes it from other more prevalent applications of generative AI such as image and video generation, where an incorrectly generated pixel might have no perceivable impact on the final artifact, and large image and video datasets are widely available. In comparison, slight errors in design representations can render a design infeasible. Furthermore, large datasets of engineering designs are infrequent, and often depend on access to a scientific simulation engine. These two challenges are the motivation for FORGE (Framework for Optimization and Reinforcement-driven Generative Engineering), our proposed probabilistic programming approach to generative design.
FORGE is a framework that facilitates design generation using probabilistic programming, based on a declarative specification of the design domain. FORGE is a symbolic, declarative approach that is complementary to neural-based approaches. As such, we use FORGE to both generate large datasets of designs, and to refine design distributions using reinforcement learning. We show how we can handle a mix of discrete and continuous parameters, as well as open-ended (recursive) design, wherein a finite set of primitive components can be combined in an infinite number of ways.
In Section 2, we provide a background for the work described here. In Section 3, we describe our approach in detail. In Section 4 we show results in several different domains, and in Section 5 we discuss our work in a broader context, and in Section 6, we provide some concluding remarks.
2. Background
In previous works, a precursor to FORGE was used as a data generating module. This includes Reference Cobb, Roy, Elenius, Heim, Swenson, Whittington and JhaCobb et al. (2023) and Reference Ghiglino, Elenius, Roy, Kaur, Acharya, Samplawski and CobbGhiglino et al. (2025). For these use-cases, FORGE was a vital component to both defining the structure of the design language, as well as playing the role of the data generation module.
Reference Cobb, Roy, Elenius, Heim, Swenson, Whittington and JhaCobb et al. (2023) was the first instance where we initiated the development of both a design language, and a framework for sampling designs. The goal of this work was to develop a machine learning-friendly dataset of unmanned aerial vehicle designs. While the development of such a framework was necessary to build such a dataset, the focus of the work was more on the structure of the final artifacts themselves, as well as the use of the data to benchmark ML models, such as transformers, for engineering design. Having developed this framework, we further extended it to the Webots simulator (Reference MichelMichel, O., 2004) for use in vehicle design. More recently, we extended FORGE to another flight dynamics simulator, where we are looking to design air taxis (Reference Ghiglino, Elenius, Roy, Kaur, Acharya, Samplawski and CobbGhiglino et al., 2025). Throughout the development of this framework, we have seen FORGE as the vital “design backend”, which has facilitated this line of work on generative design. This paper provides an opportunity to highlight the novelties and features that were not included in prior publications. More generally, AI for design has seen a significant focus in recent years. These works vary from replacing individual design modules with machine learning models (e.g. computational fluid dynamics, (Reference Shen and AlonsoShen & Alonso, 2025), (Reference Bonnet, Mazari, Cinnella and GallinariBonnet et al, 2022)), all the way to a direct replacement of the entire pipeline (Reference Elrefaie, Qian, Wu, Chen, Dai and AhmedElrefaie et al., 2025). In terms of design language, there are various works in the literature that have built impressive simulator-specific grammars, such as in Reference Zhao, Xu, Konaković-Luković, Hughes, Spielberg, Rus and MatusikZhao et al. (2020), or earlier examples such as Reference Schmidt, Shetty and ChaseSchmidt et al. (1998) and Reference Schmidt and CaganSchmidt & Cagan (1997). However, these grammars do not generalize to a range of simulators as is the case for FORGE. This is one of the key characteristics that distinguishes FORGE from prior work. The other aspect of FORGE that we explore in this paper is in its use as a direct design module. This relates to the use of generative AI in engineering design. This both differs from the direct use of black-box optimization techniques such as Bayesian optimization (Reference Balandat, Karrer, Jiang, Daulton, Letham, Wilson and BakshyBalandat et al., 2020) and evolutionary algorithms (Reference Pétrowski and Ben-HamidaPétrowski & Ben-Hamida, 2017), as well as the complete replacement of simulators with neural-based learned distributions (Reference Cobb, Matejek, Elenius, Roy and JhaCobb et al., 2024). Unlike these prior approaches, since FORGE is a parameterizable probabilistic program built from prior knowledge, it can take the form of both a generative model (unlike black-box optimization) and can be used as part of a design optimization pipeline via the use of reinforcement learning (RL), or maximum-likelihood estimation when good designs are already available. As such, FORGE presents itself as a unique and useful tool in the use of AI for design.
3. Approach
Figure 1 provides an overview of the FORGE framework. A domain expert creates a data model (1) and associated initial probability distributions (these are often boilerplate, and essentially uniform over the allowed ranges, to start with). FORGE samples a batch of designs from the data model and distributions (2), producing designs (3) and associated probabilities (4). The designs are fed to a simulator, which produces a reward for each design (5). FORGE uses the probabilities and rewards to update the parameters of the probability distributions, using RL (6). The dashed line denotes the components that constitute the final design generator. Green components need to be provided for each design domain. Blue components are reusable FORGE framework components. Orange (sim) is typically off-the-shelf.
FORGE architecture
Figure 1 Long description
A diagram of the FORGE architecture illustrating the process of design generation using artificial intelligence. The diagram includes several key components: expert input, data model, probability distributions, FORGE sampler, designs, simulation harness, simulation, rewards, FORGE learning, and updates. Expert input feeds into both the data model and probability distributions. The data model and probability distributions together form the design generator, which outputs designs through the FORGE sampler. These designs are then processed by the simulation harness, which includes a simulation component. The simulation harness outputs rewards, which are fed into the FORGE learning component. The FORGE learning component then provides updates to the probability distributions, completing the loop.
In the following subsections, we describe our approach from the point of view of generating designs for a given domain, starting from scratch. We break the process down into the following steps: Modelling (Section 3.1), Generation (Section 3.2), Evaluation (Section 3.3), and Training (Section 3.4). Note that, while we describe these aspects as different parts of an integrated system, they can also be used independently of one another. For example, Modelling-Generation-Evaluation can be used to generate training data for ML tasks; Evaluation alone can be used to evaluate manually generated designs; and Training can be coupled with a different generative model (e.g., a neural network).
3.1. Modelling
The first step in our approach is to create a data model, or schema, for our designs. We use Python, and the dataclasses and typing packages as fundamental building blocks. We use UAV design as our running example. The data model contains one top-level class (e.g., UAV) and usually contains many other classes to represent components (e.g., Propeller, Wing). It should be constructed such that any conceivable design is a valid instance of the top-level class. It is important to not inject too much designer bias and try to limit the data model to only “good” designs at this stage. The following features are implemented by our framework: a) Primitive types: float, for continuous parameters such as chord_root; int, for discrete values or counts; str, used for enumerated types, such as battery_type (a fixed list of allowed values must be provided); and bool, to indicate presence or absence of some feature, b) Structured types: List (with a type parameter), usually to indicate an unknown number of a sub-component, e.g., prop_arms; and dataclasses, with fields containing values of the previously mentioned types, or other dataclasses. We also allow for abstract classes and subclasses, which can be used to indicate a choice of wholly different types of sub-structures (e.g., VerticalTail vs. HorizontalTail). Many other types (e.g., Union, Tuple, Optional) could easily be added to the framework if needed. Figure 2 shows an example of (part of) a data model for UAVs. The top-level class Aircraft may contain different types of wings (MainWing, VerticalTail, HorizontalTail), all of which inherit from the abstract base class Wing. Finally, once the model has been created, we annotate the fields with custom metadata annotations (using the metadata attribute on dataclasses.fields). These convey information to the generative framework (described below) about the desired range of values. For float and int, we provide min and max values; for List, the min and max length of the list; for str, the set of allowed values. The ranges should not be set too restrictively; going an order of magnitude outside of “known good” value ranges is encouraged, to allow for surprising and novel designs. A key point to keep in mind when writing the data model is that it maps out the space of design choices, and therefore the space of possible designs.
A partial data model for Aircraft designs
Part of a data model for UAVs (recursive)
3.1.1. Recursive data models
The data model can be recursive, and this opens the possibility of a very open-ended kind of novelty and creativity. To create a recursive model, we must use an abstract class that has both non-recursive “leaf node” subclasses, and recursive subclasses that refer back up to the abstract class. As an example of how this can be used, consider a UAV that has a body that is built out of tubular segments, which can be connected using different kinds of joints (2-way, 3-way, etc). The end of a tube can also be connected to a component like a propeller or a wing. This can be modelled using a structure like in Figure 3. Here, we have a class ConnectedHub3_Sym (which represents the central hub of the UAV; this variant has 3 symmetric segments going out from it). This class has a field mainSegment of type MainSegment. The latter is an abstract class. It has non-recursive subclasses like PropArm and WingArm, and recursive subclasses like CrossSegment and RotatedArmSegment. These, in turn, have fields with type MainSegment. This recursive structure allows us to build arbitrarily deep tree structures (of course, at some point, we must pick non-recursive classes to finish our design). Figure 4 shows some examples of the diverse range of designs that can be generated with this approach.
Creative designs, made possible using a recursive data model (picture adapted from Reference Cobb, Roy, Elenius, Heim, Swenson, Whittington and JhaCobb et al. (2023))
3.1.2. Probability distributions
The next step is to assign a probability distribution to each choice point, and to give it initial parameter values (prior). The distribution types must be chosen to match the type of choice it represents (discrete vs continuous, infinite vs limited range). We use PyTorch’s built-in distribution classes, and we use PyTorch’s Parameter class for the parameters of the distributions (e.g., the mean and variance of a Gaussian distribution). We discuss learning these parameters in the Section 3.4. The initialization of these parameters should not be too restrictive as we do not want to bias the generator initially.
To specify the distributions and their initial values, we create a subclass of an abstract class ProbTreeModel that is built into FORGE, and implement several abstract methods, whose role is to initialize the distribution parameters, and to return the full distribution objects, for each choice point in our data model (shown as Probability Distributions in Figure 1). We create one distribution for each class-field combination, as well as for the choice of subclass of an abstract class, and for the length of lists. In practice, this step is typically boilerplate, and we used the same model for all our domains, summarized in Table 1. However, it is possible to customize this, even down to having different distributions for specific classes or fields.
Types of distributions and parameters
The subclass of ProbTreeModel that we create is also (indirectly) a subclass of PyTorch’s Module class. The use of PyTorch parameter and module classes ensures FORGE can utilize the automatic differentiation backend to enable training, as we will discuss in Section 3.4.
3.2. Generation
The next step is to generate designs from our data model and the associated probability distributions. Generation works in a top-down manner, starting from the root class (e.g., UAV), and at each choice point, sampling from the associated distribution. The code that executes this process (FORGE Sampler in Figure 1) is essentially a meta-program, or interpreter, which processes the data model using Python’s reflection capabilities. Along the way, we also calculate the log-likelihood of each design (which is simply the sum of the log-likelihoods of each individual sampling operation, as these are assumed to all be independent). The log-likelihoods are used by the training algorithm described in Section 3.4.
The domain-specific generator can be considered a probabilistic program, but no actual programming is needed from the designer. The meta-program of FORGE takes care of interpreting the model and generating designs. By this stage, FORGE can now generate random designs (using the forward method of our model class) and get log-likelihoods for each generated design for the given domain.
3.3. Evaluation
The next step is to be able to evaluate our generated designs. We assume that we already have the necessary simulator(s), such as a flight dynamics simulator, to evaluate designs in some input format (but usually not directly in the data model we have been using). We also assume that our simulators return some output data that we can use to score our designs. Thus, the next step is to write a simulation harness that a) translates from the high-level designs produced by the FORGE sampler, to the low-level format that is natively consumed by the simulator, b) invokes the simulator(s) from our Python framework, and c) processes the results into a usable form (such as a pandas DataFrame), and returns a reward based on attributes in the result objects. The reward function should incorporate all the requirements on design and give higher scores to better (simulated) performance, according to the relevant design metrics.
3.4. Training
While we can stop here and use our derived data model to generate a large dataset for machine learning tasks, we can also directly learn a more suitable generator by optimizing the distributional parameters. The generator produces random designs and their log-likelihoods, and a way to compute rewards for designs. Now, we can finally connect the pieces and introduce Reinforcement Learning (RL). The final output of this step is simply a better tuned generator (i.e., new values for the parameters in the generator). Subsequently, we can generate as many designs as we want from the improved generator.
At a high level, the training algorithm proceeds as follows. We generate a batch of random designs and evaluate them in our black-box simulator. The evaluation gives a reward to each design, according to our expert-designed reward function. Then, we compute a loss and update the probabilistic parameters of the generator using gradient descent, such that higher reward designs become more likely. This process repeats over many episodes, until we have learned a generator for designs that fit the intended purpose. Note that our specific use of RL is what is known as a “multi-armed bandit” problem. In other words, generating a complete design is done as one atomic action, and each episode consists of exactly one such action. Specifically, we use a bandit version of PPO (Reference Schulman, Wolski, Dhariwal, Radford and KlimovSchulman et al., 2017).
Our framework can also complete partial designs, which can be of great use for system designers, when part of the design is constrained in some way. This completion of partial designs allows design requirements to be strictly enforced, while maintaining the ability of the generator to sample from, and hence explore, the remainder of the design space. The generator then samples only the missing parameters, and the training algorithm trains those parameters, while holding the rest of the parameters fixed.
4. Results
Here, we show results using our approach in several domains. Our first domain uses Webots (Reference MichelMichel, O., 2004), an open-source mobile robot simulation software developed by Cyberbotics Ltd. Webots is a relatively lightweight system, suitable for experimentation and algorithm development. We consider a domain of highly unconventional land vehicles, and construct an obstacle course consisting of a racetrack with multiple obstacles in it of varying height and size. The goal is to design a vehicle that can finish this track successfully and quickly. We set the reward function to be the sum of rewards obtained by crossing each of the obstacles discounted by the total time taken to complete the course. The data model provides several avenues for novelty in the design of the vehicle, including: (a) allowing chassis to be non-rigid and comprising several components connected via flexible joints, wheels that can be of different sizes and shapes such as cylindrical, spherical or square, and varying offsets of wheels. Figure 5 shows the environment and an example of a generated vehicle design.
The Webots test environment, and an example of a generated vehicle design
IsaacSim parallel batch evaluation
Figure 6 Long description
Panel A: A grid of rectangular panels with lines and small objects on each panel.
Subsequently, we adapted the same data model to use IsaacSim (NVIDIA), a high-fidelity 3D modeling environment and physics simulator. This has the advantage that we can run an entire batch (e.g., 32 designs) simultaneously, with GPU-acceleration (see Figure 6). This alleviates the CPU bottleneck of running multiple parallel complete instances of a simulator like Webots.
We also looked at the domain of rocket design. Here, we made use of the RocketPy simulator (Reference Ceotto, Schmitt, Alves, Pezente and CarmoCeotto et al., 2021). We ran different experiments, with reward functions to optimize for max apogee, precision landing, and target altitude. Figure 7 shows an example of generated rocket design.
A rocket design from RocketPy (left), and sample output data from the simulator (right)
Examples from generated air taxis using FORGE for machine learning training data
In our most recent work (Reference Ghiglino, Elenius, Roy, Kaur, Acharya, Samplawski and CobbGhiglino et al., 2025), we connected FORGE to an air taxi flight dynamics simulation pipeline based on Suave (Reference MacDonald, Clarke, Botero, Vegh and AlonsoMacDonald et al., 2017). We use lift coefficient over drag coefficient as the reward. Figure 8 shows some of the designs produced by FORGE. These designs show the variations in both the individual components of the aircraft, such as wing dimensions, and the global properties of the design, such as the number of wings and the configuration of propulsion systems.
Summary of domains. Sim times on AMD 9950X CPU, Nvidia RTX 5090 GPU, 96GB RAM
Table 2 summarizes some quantitative metrics of the four domains, and Figure 9 shows the reward curves from the RL training process on each domain. We use the same hyperparameters in all cases: 500 episodes, 4 epochs, batch size 32, clip 0.2, learning rate 0.01 with Adam optimizer. We compare against a baseline of random sampling from the initial (untrained) program.
Rewards per episode: Mean (solid) and standard deviation (shaded)
One issue that comes up when submitting randomly sampled designs to a simulator is that simulators have many hard constraints, and these are often difficult to know a priori. If the constraints are known, they can be included in the reward function, but there is a much simpler, pragmatic solution that works well in practice: In our sim harness, when a design fails (often signalled by the simulator raising an exception, or returning an error code), we give it a zero reward. Then, it is up to the RL process to avoid those situations. For example, RocketPy had a particularly high initial failure rate of about 60%. After 150 episodes, less than 5% of generated designs fail to satisfy all the hard constraints.
In Figure 10, we show some of the probability distributions learned through this method, from the Suave domain. On the first row, we see examples of continuous distributions (note that the x axis is normalized to map the full [min,max] range of each field to [0,1]). In the second row of the figure, we see examples of list-length distributions. The third row shows some of the subclass-choice distributions. The learned distributions describe several interesting trends. Firstly, the distribution of the wing x position is biased towards the upper end, indicating that the wings should be on average towards the back of the input range for aircraft stability. Secondly, the optimal span is found to be around 60% of the maximum span. Increasing the span decreases the drag induced by the lift of the aircraft but increases the drag due to air resistance (known as profile drag). At lower speeds, induced drag dominates, so to maximize the lift and minimize the drag, the algorithm is incentivized to increase the span until the skin friction drag or structural constraints dominate. Finally, the number of inner wing cross-sections is aligned with the numbers of panels on most aircraft, with most general aviation aircraft being one panel (0 inner cross sections), and commercial aircraft being two or three panel wings (1 or 2 inner cross sections). The full model includes about 100 distributions (and each distribution has either one or two parameters). Note that the entire learned model is fully interpretable (unlike neural models), and we can immediately see the meaning of each parameter and distribution.
Examples of learned probability distributions for continuous, list-length, and subclass-choice parameters, in the UAV domain
Figure 10 Long description
Panel A: A line graph titled MainWing_origin_x_rel shows the relationship between an unspecified variable on the x-axis and the relative origin position on the x-axis. The line starts low and increases steeply towards the right. Panel B: A line graph titled MainWing_origin_z_rel displays the relationship between an unspecified variable on the x-axis and the relative origin position on the z-axis. The line starts high and decreases steeply towards the right. Panel C: A line graph titled MainWing_span_proj illustrates the relationship between an unspecified variable on the x-axis and the projected span of the main wing. The line forms a parabolic shape, peaking in the middle. Panel D: A bar graph titled MainWing_inner_cross_sections shows the probability distribution of different cross-sectional types. The bars are labeled from 0 to 10, with probabilities ranging from 0.02 to 0.37. Panel E: A bar graph titled MainWing_prop_arms displays the probability distribution of different propulsion arm configurations. The bars are labeled from 0 to 10, with probabilities ranging from 0.01 to 0.19. Panel F: A bar graph titled RandomFuselage_inner_cross_sections shows the probability distribution of different inner cross-sectional types for the fuselage. The bars are labeled from 0 to 10, with probabilities ranging from 0.01 to 0.33. Panel G: A bar graph titled PropArm shows the probability distribution of different propeller arm types. The types are LiftRotorArmSymmetric, LiftRotorArmAsymmetric, and ForwardPropArm, with probabilities of 0.45, 0.22, and 0.34, respectively. Panel H: A bar graph titled Fuselage shows the probability distribution of different fuselage types. The types are RandomFuselage and PowerSeriesFuselage, with probabilities of 0.51 and 0.49, respectively.
5. Discussion
In this work, we have adopted a single-step (multi-armed bandit) formulation of reinforcement learning in which complete designs are generated and evaluated as atomic actions. While computationally convenient and well aligned with our probabilistic program structure, this approach does not fully leverage the inherently sequential nature of design construction. A natural extension of FORGE is to support multi-step policies, particularly through its existing partial design instantiation mechanism, enabling incremental decision-making, hierarchical refinement, and improved credit assignment across complex design choices. Another important direction concerns the measurement of diversity in generated designs. Because FORGE defines an explicit probabilistic program, artifacts are sampled from a learned distribution rather than an implicit neural representation, creating an opportunity to define diversity in distributional terms (e.g., entropy, structural coverage, or parameter spread) that can be interpreted and validated by domain experts. Finally, although we have focused on reinforcement learning to optimize the parameters of the probabilistic program, alternative strategies are possible. In preliminary experiments, Bayesian optimization proved impractical due to the high dimensionality of the parameter spaces considered (up to 186 trainable parameters), whereas gradient-based RL scaled naturally and leveraged batch simulation effectively.
6. Conclusion
In this paper we have introduced FORGE, a probabilistic framework for engineering designs. We have focused on the illustrative examples of vehicle designs, whereby we have shown how one adapts FORGE to varying design domains. We then demonstrated how the framework works, in setting up the data model, and the ability to generate a large variety of designs thereafter. At this point, one can use these generated designs along with the domain-specific simulator to build a large dataset for machine learning tasks as we have done in our prior works. Additionally, we have shown how we can directly optimize over the data model’s probabilistic distributional parameters to achieve highly performant designs. This optimized design distribution is a more interpretable generative model than neural-based generative distributions. Exploring multi-step learning formulations, principled diversity metrics, additional optimization strategies, and applications to larger-scale, real-world design problems, including integration with existing CAD workflows and complementary AI tools, represents promising future work.
Acknowledgements
This material is based upon work supported by the United States Air Force and DARPA under Contract No. FA8750-23-C-0519 and HR0011-24-9-0424, and the U.S. Army Research Laboratory under Cooperative Research Agreement W911NF-17-2-0196. Any opinions, findings and conclusions or recommendations expressed in this material are those of the author(s) and do not necessarily reflect the views of the United States Air Force, DARPA, the U.S. Army Research Laboratory, or the United States Government.
Open access
