📐TorchArc

Declarative YAML Network Architectures

Code: slm_lab/agent/net/torcharc_net.py · Library: torcharc

TorchArc builds neural networks from declarative YAML specs via the torcharc library. Instead of hardcoded network classes like MLPNet or ConvNet, you define the exact PyTorch modules and dataflow graph in YAML. All benchmark_arc/ specs use TorchArc (v5.1+).

YAML Structure

A TorchArc spec has two parts: modules (what layers to build) and graph (how data flows through them).

net:
  type: TorchArcNet
  shared: false
  arc:
    modules:
      body:
        Sequential:
          - LazyLinear: {out_features: 256}
          - ReLU:
          - LazyLinear: {out_features: 256}
          - ReLU:
    graph:
      input: x
      modules:
        body: [x]
      output: body
  hid_layers_activation: relu
  clip_grad_val: 0.5
  optim_spec:
    name: Adam
    lr: 3.0e-4
  gpu: auto

• modules — named groups of PyTorch modules. Each key (e.g. body) becomes a callable sub-network. Use any torch.nn module by name.

• graph — defines the input name, which modules receive which inputs, and which module produces the final output.

• LazyLinear / LazyConv2d — input dimensions are inferred automatically, so you only specify output sizes.

YAML Anchors for Compact Specs

Real benchmark specs use YAML anchors (&name / *name) to define a base config once and reuse it across environments. This is the key to keeping multi-environment specs DRY.

Define once, reuse everywhere

# Define the base architecture once
_ppo_mujoco_arc: &ppo_mujoco_arc
  modules:
    body:
      Sequential:
        - LazyLinear: {out_features: 256}
        - Tanh:
        - LazyLinear: {out_features: 256}
        - Tanh:
  graph:
    input: x
    modules:
      body: [x]
    output: body

# Define the base net config once
_ppo_mujoco_net: &ppo_mujoco_net
  type: TorchArcNet
  shared: false
  arc: *ppo_mujoco_arc          # ← reference the architecture
  hid_layers_activation: tanh
  init_fn: orthogonal_
  clip_grad_val: 0.5
  use_same_optim: true
  optim_spec:
    name: AdamW
    lr: 3.0e-4
  gpu: auto

Merge and override

Use <<: *anchor to inherit all fields, then override only what differs. Here's the PPO Atari spec — three lambda variants defined by changing a single line each:

# --- Shared anchors (defined once at top of file) ---

_ppo_atari_algorithm: &ppo_atari_algorithm
  name: PPO
  gamma: 0.99
  lam: 0.95
  clip_eps_spec: {name: no_decay, start_val: 0.1, end_val: 0.1}
  entropy_coef_spec: {name: no_decay, start_val: 0.01, end_val: 0.01}
  val_loss_coef: 0.5
  time_horizon: 128
  minibatch_size: 256
  training_epoch: 4

_ppo_atari_net: &ppo_atari_net
  type: TorchArcNet
  shared: true
  arc: *ppo_atari_arc
  hid_layers_activation: relu
  init_fn: orthogonal_
  optim_spec: {name: AdamW, lr: 2.5e-4, eps: 1.0e-5}
  lr_scheduler_spec: {name: LinearToZero, frame: 1.0e+7}
  normalize: true
  gpu: auto

_ppo_atari_env: &ppo_atari_env
  name: "${env}"
  num_envs: 16
  max_frame: 1.0e+7

# --- Specs: only overrides visible ---

ppo_atari_arc:                     # lam=0.95 (default)
  agent:
    algorithm: *ppo_atari_algorithm
    net: *ppo_atari_net
  env: *ppo_atari_env

ppo_atari_lam85_arc:               # lam=0.85 — one line changed
  agent:
    algorithm:
      <<: *ppo_atari_algorithm
      lam: 0.85                    # ← only difference
    net: *ppo_atari_net
  env: *ppo_atari_env

ppo_atari_lam70_arc:               # lam=0.70 — one line changed
  agent:
    algorithm:
      <<: *ppo_atari_algorithm
      lam: 0.70                    # ← only difference
    net: *ppo_atari_net
  env: *ppo_atari_env

Why this matters: Three complete Atari specs differ by a single line each. Without anchors, each would repeat ~40 lines of identical config. When scanning the file, the differences jump out immediately.

The same pattern works for per-environment overrides in MuJoCo — only the tuned hyperparameters are visible:

ppo_ant_arc:
  agent:
    algorithm:
      <<: *ppo_mujoco_algorithm
      gamma: 0.988                # ← tuned for Ant
      lam: 0.928
    net:
      <<: *ppo_mujoco_net
      optim_spec:
        name: AdamW
        lr: 1.5e-4                # ← lower lr for Ant

ppo_hopper_arc:
  agent:
    algorithm:
      <<: *ppo_mujoco_algorithm
      gamma: 0.991                # ← tuned for Hopper
    net:
      <<: *ppo_mujoco_net
      lr_scheduler_spec:          # ← Hopper needs lr decay
        name: LinearToZero
        frame: "${max_frame}"

Before vs. after

Without anchors — repetitive, hard to diff

ppo_ant_arc:
  agent:
    net:
      type: TorchArcNet
      shared: false
      arc:
        modules:
          body:
            Sequential:
              - LazyLinear: {out_features: 256}
              - Tanh:
              - LazyLinear: {out_features: 256}
              - Tanh:
        graph:
          input: x
          modules:
            body: [x]
          output: body
      hid_layers_activation: tanh
      init_fn: orthogonal_
      clip_grad_val: 0.5
      use_same_optim: true
      optim_spec:
        name: AdamW
        lr: 1.5e-4        # ← the only actual difference
      gpu: auto

ppo_hopper_arc:
  agent:
    net:
      type: TorchArcNet
      shared: false
      arc:
        modules:
          body:
            Sequential:
              - LazyLinear: {out_features: 256}
              - Tanh:
              - LazyLinear: {out_features: 256}
              - Tanh:
        graph:
          input: x
          modules:
            body: [x]
          output: body
      hid_layers_activation: tanh
      init_fn: orthogonal_
      clip_grad_val: 0.5
      use_same_optim: true
      optim_spec:
        name: AdamW
        lr: 3.0e-4
      lr_scheduler_spec:   # ← the only actual difference
        name: LinearToZero
        frame: "${max_frame}"
      gpu: auto

With anchors — only overrides are visible:

ppo_ant_arc:
  agent:
    net:
      <<: *ppo_mujoco_net
      optim_spec: {name: AdamW, lr: 1.5e-4}

ppo_hopper_arc:
  agent:
    net:
      <<: *ppo_mujoco_net
      lr_scheduler_spec: {name: LinearToZero, frame: "${max_frame}"}

Atari (Conv) Architecture

TorchArc handles convolutional networks the same way—list the modules explicitly:

_ppo_atari_arc: &ppo_atari_arc
  modules:
    body:
      Sequential:
        - LazyConv2d: {out_channels: 32, kernel_size: 8, stride: 4, padding: 0}
        - ReLU:
        - LazyConv2d: {out_channels: 64, kernel_size: 4, stride: 2, padding: 0}
        - ReLU:
        - LazyConv2d: {out_channels: 64, kernel_size: 3, stride: 1, padding: 0}
        - ReLU:
        - Flatten:
        - LazyLinear: {out_features: 512}
        - ReLU:
  graph:
    input: x
    modules:
      body: [x]
    output: body

This is the Nature CNN architecture. With TorchArc, the conv layers, flatten, and FC head are all visible in one place.

Adding Normalization Layers

Insert any torch.nn module into the Sequential list. For example, adding BatchNorm:

Sequential:
  - LazyLinear: {out_features: 256}
  - LazyBatchNorm1d: {}
  - ReLU:
  - LazyLinear: {out_features: 256}
  - LazyBatchNorm1d: {}
  - ReLU:

Or LayerNorm:

Sequential:
  - LazyLinear: {out_features: 256}
  - LayerNorm: {normalized_shape: 256}
  - ReLU:
  - LazyLinear: {out_features: 256}
  - LayerNorm: {normalized_shape: 256}
  - ReLU:

With TorchArc, adding normalization is a one-line insertion per layer—no code changes needed.

Old vs. New

	MLPNet (old)	TorchArcNet (new)
Definition	hid_layers: [256, 256]	Explicit YAML modules + graph
Reuse	Copy-paste across specs	YAML anchors (& / *)
Flexibility	Fixed patterns (MLP, Conv, RNN)	Any torch.nn module
Transparency	Layers are implicit	Every layer is visible
Normalization	batch_norm: true flag	Insert BatchNorm1d, LayerNorm, etc. directly

All benchmark_arc/ specs use TorchArcNet. The old MLPNet/ConvNet types still work for backward compatibility, but TorchArc is the recommended approach for new specs.

What's Next

• Benchmark Specs — see TorchArc specs in action

• torcharc library — full documentation for the YAML spec format

• Browse slm_lab/spec/benchmark_arc/ for all benchmark TorchArc specs