Enhance ViT policy with configurable trunk and update TF32 precision settings

[daveey]

No description provided.

[daveey]

Warning

This pull request is not mergeable via GitHub because a downstack PR is open. Once all requirements are satisfied, merge this PR as a stack on Graphite.
Learn more

• Add script to create WandB table with run metadata #4172
• Enhance ViT policy with configurable trunk and update TF32 precision settings #4148 👈 (View in Graphite)
• Add Muesli algorithm implementation with dynamics unrolling and CMPO policy optimization #4146
• fix: SimulationAgent.inventory reads only from agent's own position #4239
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.

[graphite-app]

NameError: Undefined variable reference

The code references trunk_num_resnet_layers and trunk_use_layer_norm in the components list default value, but these variables are not in scope at class definition time. Class attributes defined above (lines 165, 167) cannot be referenced by name in other class attribute default values.

This will cause an immediate NameError when the class is defined:

NameError: name 'trunk_num_resnet_layers' is not defined

Fix: Either use hardcoded values in the default or use _latent_dim pattern with underscore prefix, or move this logic to __post_init__ or the make_policy method where instance attributes are accessible.

Suggested change

	num_layers=trunk_num_resnet_layers,
	pattern="A", # Axon blocks provide residual-like connections
	post_norm=trunk_use_layer_norm,
	num_layers=4, # Default value
	pattern="A", # Axon blocks provide residual-like connections
	post_norm=True, # Default value

Spotted by Graphite Agent



Is this helpful? React 👍 or 👎 to let us know.

[openhands-ai]

Looks like there are a few issues preventing this PR from being merged!

• GitHub Actions are failing:
  • Test and Benchmark

If you'd like me to help, just leave a comment, like

@OpenHands please fix the failing actions on PR #4148 at branch `daveey-vitres`

Feel free to include any additional details that might help me get this PR into a better state.

_{^{You can manage your notification settings}}