neural os - Moses Yebei

The Problem

Training agents to interact with a real operating system is slow, serial, and hard to parallelize: every rollout needs a live desktop, and failures are expensive. NeuralOS replaces the OS with a neural model that predicts the next screen frame from mouse, click, and keyboard events, so policies can be trained and evaluated against a GPU-resident, differentiable environment instead of a physical machine.

What This Does

The repo is a fork of yuntian-group/neural-os and contains three training pipelines. autoencoder/ trains a VQ/KL autoencoder that compresses 512×384×3 frames into 64×48×16 latents (an 8x reduction). latent_diffusion/ is the CompVis latent-diffusion framework, adapted to render those latents back into images. computer/ trains the main model: an RNN tracking OS state plus a diffusion renderer, driven by configs in computer/configs/. data/ generates and aggregates the Ubuntu XFCE training recordings.

How It Is Wired

Execution starts at three real entry points. process_video in data/data_processing/main.py:65 aggregates recordings; its traced path reaches the filesystem in one hop via Image.open inside compute_distance. main in autoencoder/visualize_reconstruction_kl_f16.py:108 reaches 64 functions and leaves the process via os.makedirs. The training entry points are autoencoder/main.py and computer/main.py, both driven by YAML configs.

Three modules are hubs with wide blast radius: computer/util (32 modules depend on it), data/data_processing/datasets.py (31), and latent_diffusion/ldm/util (20). The core of the system is latent_diffusion/ldm/models/diffusion/ddpm.py: 72 functions, 4 classes, called from 11 files, 1772 lines — a change there ripples everywhere. The most-called functions are exists (27 call sites), sample (20), and instantiate_from_config (11). latent_diffusion/ldm/models/diffusion/ddim.py (16 callers) and latent_diffusion/ldm/modules/attention.py (20 callers) are also widely referenced. No circular dependencies were found in the 150-module import graph.

How To Use It

The README documents the full pipeline; commands are used verbatim:

# 1. Data generation (build the collection environment)
cd data/data_collection
docker build -t synthetic_data_generator .

# 2. Data processing
cd data/data_processing
python main.py

# 3. Autoencoder training
cd autoencoder/
python main.py --config config_kl4_lr4.5e6_load_acc1_512_384_mar10_keyboard_init_16_contmar15_acc1_cont1e6.yaml

# 4. Image preprocessing to latent WebDataset tars
python preprocess_dataset.py

# 5. Main model training (8 GPUs, FSDP)
cd computer/
python main.py --config configs/fb_computecanada_challengingandsample_pretrainrnn_balanced_lr5e6_contbest_samplercover_newd_contfreezernn_newnewd_origunet_nospatial_online_x0_joint_onlineonly_7.yaml

Training the autoencoder takes roughly a week on one H200. Config paths must be edited to point at your data; multi-stage training is controlled by flags like "pretrain": true and "freezernn": true in the YAML.

Real-World Use

A practical fit: training a GUI agent policy in the loop. Generate 40K random and agent-driven Ubuntu sessions via the Docker container, encode them to latents, train the autoencoder, then train the RNN+diffusion model to predict frames. The trained model becomes a fast, parallelizable surrogate for evaluation — useful for RL and for testing agents without provisioning desktops.

Code Health & Issues

Static analysis (measured, not opinion) found 99 issues: 32 high, 65 medium, 2 low. Key findings:

High - deep nesting (x37) — data/data_processing/datasets.py, latent_diffusion/ldm/models/diffusion/ddpm.py — max indentation depth 10; flatten with guard clauses.
High - oversized files — ddpm.py (1772 lines), latent_diffusion/ldm/modules/diffusionmodules/model.py — split by responsibility.
High - hub modules — computer/util.py, datasets.py, latent_diffusion/ldm/util.py — churn here is high-blast-radius.
High - duplicated code — 2079 repeated 6-line blocks across 136 files; extract shared helpers.
Medium - broad exception handling (x2) — bare except in datasets.py, ddpm.py.
Medium - file handles without context managers (x14) — use with open(...).

SDLC observations: High — no LICENSE file (all rights reserved by default, blocking reuse). High — no CI workflow for 150 source files. High — data/data_collection/Dockerfile is a deployable artifact with no build gate. Medium — no Dependabot/Renovate; base image ubuntu:20.04 unpinned; no non-root USER in the Dockerfile; test ratio is 2 test files to 150 source files (0.013), so the suite cannot cover this codebase.

The Bottom Line

This is a serious research codebase that reproduces a published pipeline end-to-end, but it is a fork of a fork: heavy config sprawl (287 YAML files, many near-duplicate experiment configs), no CI, and no license. Use it if you need a working neural-OS simulator for agent training and can tolerate a research-grade build; harden it before any production or client deployment.

What the analyser found

Deployment readiness

5/7

Container image
CI pipeline
Lockfile committed
Test suite
README
License
No committed secrets

Composition

1,065 files

YAML287
Python150
Shell55
Markdown4
JSON1

Docker

Architecture

Top-level areas of the codebase, sized by module count. Arrows show how many imports cross from one area into another.

graph TD g0["computer 56 modules"] g1["latent_diffusion 34 modules"] g2["autoencoder 33 modules"] g3["data 22 modules"] g4["(root) 5 modules"] g2 -->|28| g0 g2 -->|27| g3 g2 -->|23| g1 g0 -->|16| g1 g0 -->|5| g3 g3 -->|2| g1

Module dependencies

The 10 most depended-upon modules of 150, from static import analysis. Red outline marks a module in an import cycle.

graph LR n0_computer_util["util"] n1_ata_data_processing_datasets["datasets"] n2_latent_diffusion_ldm_util["util"] n3_on_ldm_models_diffusion_ddpm["ddpm"] n4_usion_ldm_models_autoencoder["autoencoder"] n5_odules_diffusionmodules_util["util"] n6_on_ldm_models_diffusion_ddim["ddim"] n7_fusion_ldm_modules_attention["attention"] n8_lection_synthetic_mouse_path["synthetic_mouse_path"] n9__distributions_distributions["distributions"] n0_computer_util --> n1_ata_data_processing_datasets n0_computer_util --> n2_latent_diffusion_ldm_util n0_computer_util --> n3_on_ldm_models_diffusion_ddpm n1_ata_data_processing_datasets --> n2_latent_diffusion_ldm_util n4_usion_ldm_models_autoencoder --> n9__distributions_distributions n4_usion_ldm_models_autoencoder --> n2_latent_diffusion_ldm_util n6_on_ldm_models_diffusion_ddim --> n5_odules_diffusionmodules_util n3_on_ldm_models_diffusion_ddpm --> n2_latent_diffusion_ldm_util n3_on_ldm_models_diffusion_ddpm --> n9__distributions_distributions n3_on_ldm_models_diffusion_ddpm --> n4_usion_ldm_models_autoencoder n3_on_ldm_models_diffusion_ddpm --> n5_odules_diffusionmodules_util n3_on_ldm_models_diffusion_ddpm --> n6_on_ldm_models_diffusion_ddim n7_fusion_ldm_modules_attention --> n5_odules_diffusionmodules_util n5_odules_diffusionmodules_util --> n2_latent_diffusion_ldm_util

Dependencies

9 declared dependencies, none a major version behind. Resolved 9 of 9 against the registry.

Findings

MediumNo CI/CD pipeline detected - no automated build/test gate .github/ or CI config
LowDependencies declared without a lockfile - non-reproducible builds data/data_collection/requirements.txt

What would stop me shipping this

Ranked by severity × confidence × production reach. Reach is the honest discriminator across a collection that is mostly other people's code: the same finding matters more in something that ships.

high3

medium4

Add a LICENSE; redistribution rights are undefinedhigh

no licence file at the repository root

Fix: Add MIT or Apache-2.0 at the repository root.

Add a workflow that builds and tests this repositoryhigh

150 source files, no CI configuration

Every change merges with nobody having run the build once, which across this estate is 499 repositories.

Fix: Add a workflow running the project build and test command on push and pull_request.

Add a build gate for the deployable artifacts herehigh

data/data_collection/Dockerfile

deployment artifacts with no workflow

Something in the repo is meant to reach a server and no automated step has ever validated it, which is the state of 125 repositories here.

Fix: Add a workflow that at least builds the image and validates the manifests.

Enable Dependabot or Renovatemedium

2 manifest(s), no update bot configured

Without a bot a published advisory sits unpatched until someone audits by hand, which across 1,322 repositories means never.

Fix: Commit .github/dependabot.yml covering the repo ecosystems plus github-actions.

Pin the container base image by digestmedium

data/data_collection/Dockerfile

ubuntu:20.04

An untagged or mutable base means today's build and last month's contain different libc and a different CVE set, with no record of which shipped.

Fix: Use image:tag@sha256:<digest> and enable Dependabot's docker ecosystem.

Add a non-root USER to the imagemedium

data/data_collection/Dockerfile

CMD or ENTRYPOINT with no USER directive

A process running as root in the container is root against every mounted volume, and it turns any container escape or writable-mount mistake from a contained problem into a host one.

Fix: Create an unprivileged user, chown what it needs, and end the Dockerfile with USER.

Expand the test suite; it cannot cover this codebasemedium2 occurrences

2 test files against 150 source files (ratio 0.013)

A reviewer reads "has tests" and assumes a safety net that covers two files out of eighty-five, which is how untested paths ship behind a green badge.

Fix: Add tests for the highest fan-in modules first.

Checked deterministically against the repository tree and a bounded set of its files: committed credentials, unpinned actions and base images, missing lockfiles and update bots, workflows that discard failures, published advisories against the declared dependencies, runtime configuration, licensing and notebook reproducibility. No language model is involved in this section.