Your Robot Thinks in UML Now — And That's a Very Good Thing

Borrowing class diagrams and activity diagrams from software engineering, OOWM gives embodied AI agents an explicit, structured world model instead of the fuzzy linear text that's been holding them back.

~6 min read · arXiv / cs.AI · 2026-04-14 · Architecture

TL;DR Chain-of-Thought prompting asks LLMs to reason in natural language, but free-form text is a terrible data structure for representing a physical room full of objects, states, and causal dependencies. OOWM replaces that text soup with UML — the same Class Diagrams and Activity Diagrams software engineers use to model systems — and trains a vision-language model to generate them. The result is a robot brain that can distinguish "the cup is dirty" from "the cup is fragile" and then produce an executable plan to deal with it.

There's a quiet irony in how we've been teaching robots to reason. We give them billion-parameter models with access to the sum of human knowledge, then ask them to represent a kitchen — its drawers, its states, its causal logic — as a paragraph of English. A language designed for poetry and persuasion, pressed into service as a data structure for robotics.

The team behind OOWM (Object-Oriented World Modeling) decided to try something more principled: borrow the tools software engineers have used for decades to model complex stateful systems.

💡 The Core Idea A physical environment is, structurally, a lot like a software system — objects with attributes, states that change, actions with preconditions and postconditions. UML was invented to model exactly that. OOWM asks: what if the robot's "thinking" was a UML diagram instead of a sentence?

The Problem with "Thinking Out Loud"

Chain-of-Thought (CoT) prompting works by encouraging a model to write out its reasoning step by step before giving a final answer. For math problems and logic puzzles, this is enormously effective. For embodied tasks — "clean this room" — it falls apart in three specific ways the authors identify.

First, natural language blurs the distinction between an object's static attributes (a vase is fragile) and its dynamic state (the vase is currently on the floor). These are fundamentally different kinds of information that require different handling in a plan, but plain text treats them identically.

Second, verifying logical consistency across a long-horizon plan is nearly impossible in prose. Did step 7 assume the table was clear, when step 4 put something on it? A paragraph won't catch that contradiction; a formal state machine would.

Third, text requires an extra translation step to become something a robot can actually execute. You've reasoned about cleaning the room in English; now someone has to convert that into a control policy.

The OOWM Solution: Two Diagrams, One World Model

OOWM formally defines the world model as a symbolic tuple:

$$W = \langle S, T \rangle$$

In plain English: the world model is two things — a snapshot of the current state of the environment (S), and a description of how actions change that state (T: S × A → S').

These two components are materialized as two kinds of UML diagram, both serialized as PlantUML code that the model actually generates as text:

State Abstraction (G_state) → UML Class Diagram. When the agent looks at a scene, instead of writing "there's a dirty cup next to a book," it generates a class diagram. Objects become classes. Attributes (color, material) become typed fields. Dynamic properties (dirty, open, broken) become state fields. The hierarchy — that a "ceramic mug" is a subtype of "container" — is made explicit through inheritance.

Control Policy (G_control) → UML Activity Diagram. The plan isn't a to-do list in prose. It's an activity diagram: a flowchart with decision nodes, parallel branches, and explicit control flow. "If the object is fragile, use the gentle grip subroutine; otherwise use the power grip" is not ambiguous in a diagram the way it can be in text.

flowchart LR
    A[Visual Input\nScene Image] --> B[G_state\nUML Class Diagram\nObject hierarchies & states]
    B --> C[G_control\nUML Activity Diagram\nExecutable control flow]
    C --> D[Robot Action\nExecution]
    B -.->|State informs| C

Training: Three Stages, Sparse Supervision

Teaching a model to generate valid, useful UML is non-trivial, and the authors address this with a three-stage pipeline built on top of InternVL 2.5-1B:

Stage 1 — Supervised Fine-Tuning on State Abstraction. A small subset of the MRoom-30k benchmark (1,000 samples) is fully annotated with ground-truth Class Diagrams. The model learns what a correct G_state looks like from examples.

Stage 2 — Extending to Full OOWM. The model is trained on the base planning set (approximately 29,000 samples split 80/10/10 for train/val/test) to produce both components together.

Stage 3 — Reinforcement Learning via GRPO. Here's the clever part. Ground-truth G_state annotations are expensive to create. Rather than requiring dense supervision on the intermediate reasoning structure, Stage 3 uses Group Relative Policy Optimization (GRPO) with outcome-based rewards — signals derived from whether the final plan actually works, not whether the intermediate Class Diagram was perfectly drawn. The model learns to produce better object-oriented reasoning implicitly, because better reasoning leads to better plans, and better plans are what get rewarded.

This matters practically: it means OOWM can improve its internal world modeling from the much-cheaper signal of "did the robot succeed?" rather than requiring expert annotation of every intermediate reasoning step.

What the Ablations Reveal

The authors test four configurations on MRoom-30k to isolate what's actually driving gains:

The progression from Hybrid to OOWM 2-Stage isolates the value of structuring the reasoning, not just the output. The jump from 2-Stage to 3-Stage isolates the value of outcome-based RL. The paper reports that the full three-stage OOWM significantly outperforms the unstructured text baseline across planning coherence, execution success, and structural fidelity — though the specific numeric deltas aren't available in the extracted content.

⚠️ Watch out for

• The backbone is InternVL 2.5-1B — a relatively small model. It's unclear how performance scales or whether the UML-generation overhead becomes a bottleneck at inference time on resource-constrained hardware.
• GRPO Stage 3 uses only 2,000 samples due to computational constraints. Results might differ significantly with a larger RL training set.
• MRoom-30k is the sole benchmark reported. Generalization to other embodied environments (outdoor, manipulation-heavy, dynamic) remains an open question.

Why It Matters

The deeper argument here is about the right abstraction layer for embodied AI. Neural world models — learning to predict the next frame, or encoding the environment as a latent vector — are powerful but opaque. You can't inspect them, debug them, or tell the robot "by the way, that vase is an heirloom."

OOWM points toward a different design philosophy: keep the world model explicit and symbolic, use neural networks to generate and update it from perception, and use formal structures with known properties (like UML diagrams) to make it inspectable and verifiable. A Class Diagram you can read. A latent vector you cannot.

For anyone building robotic systems that need to be auditable — in warehouses, hospitals, or home assistance — that legibility is not a minor nicety. It's a safety property.

The use of PlantUML as the serialization format is also quietly pragmatic: it means the model's "thinking" is just structured text, compatible with standard language model training pipelines, no custom tokenization or architecture changes required.

Source: OOWM: Structuring Embodied Reasoning and Planning via Object-Oriented Programmatic World Modeling Authors: Hongyu Chen, Liang Lin, Guangrun Wang Published: 2026-04-14 PDF: https://arxiv.org/pdf/2604.09580