Skip to content

Conversation

@eugenevinitsky
Copy link

@eugenevinitsky eugenevinitsky commented Jan 5, 2026

Currently the visualizer draws its config during training directly from the ini file. So, if we pass important things as CLI args, they aren't included during the visualization process. This writes the CLI args to a new config so that they are correctly pulled in during visualization.

- Add --config CLI flag to visualize.c for custom INI file
- Add _write_visualizer_config() to generate config from driver_env
- Handle special attribute mappings (action_type -> _action_type_flag)
- Pick random map from training map_dir when render_map not specified

🤖 Generated with [Claude Code](https://claude.com/claude-code)
@greptile-apps
Copy link

greptile-apps bot commented Jan 5, 2026

Greptile Summary

This PR enables the visualizer to inherit environment configuration from the training process instead of always using default values. The changes introduce a new _write_visualizer_config() helper that reads the base INI file and overrides [env] section values with attributes from the training environment. The C visualizer is extended with a --config flag to accept custom config files, and drive.py now validates that the map directory contains any .bin files rather than requiring a specific map_000.bin file.

Key Changes:

  • Added _write_visualizer_config() function in pufferlib/utils.py to generate temporary config files with training environment overrides
  • Modified render_videos() to create and pass a custom config file to the visualizer using the new --config flag
  • Extended C visualizer's eval_gif() function signature to accept a config_file parameter
  • Improved map directory validation in drive.py to check for any .bin files instead of requiring map_000.bin
  • Updated documentation to reflect the new error message format

Issues Found:

  • cfg.read() in _write_visualizer_config() (line 175) fails silently if the base INI file doesn't exist, which could result in an empty or invalid config being written

Confidence Score: 3/5

  • This PR is mostly safe but has one logic issue that could cause runtime failures
  • Score reflects one critical logic bug in config file handling that needs to be fixed - the cfg.read() call doesn't validate that the base INI file was successfully read, which could lead to silent failures and invalid visualizations. The rest of the changes are well-structured and improve flexibility, but this bug prevents a higher confidence score.
  • Pay close attention to pufferlib/utils.py - the config reading logic needs error handling

Important Files Changed

Filename Overview
pufferlib/ocean/drive/drive.py Improved error checking for map directory validation - now checks for any .bin files instead of requiring map_000.bin specifically
pufferlib/ocean/drive/visualize.c Added --config flag to allow custom config file path, enabling visualizer to inherit training environment configuration
pufferlib/utils.py Added _write_visualizer_config function and integrated it with render_videos, but missing error handling for config file reading
docs/src/data.md Updated documentation to reflect new error message for missing .bin map files

Sequence Diagram

sequenceDiagram
    participant Train as Training Process
    participant Utils as render_videos()
    participant Helper as _write_visualizer_config()
    participant ConfigFile as visualizer_config.ini
    participant Visualizer as C Visualizer (./visualize)
    
    Train->>Utils: Call render_videos(config, vecenv, ...)
    Utils->>Utils: Get env_cfg from vecenv.driver_env
    Utils->>Helper: Call _write_visualizer_config(env_cfg, temp_config_path)
    Helper->>Helper: Read base INI (pufferlib/config/ocean/drive.ini)
    Helper->>Helper: Override [env] values with env_cfg attributes
    Helper->>ConfigFile: Write modified config to temp path
    Helper-->>Utils: Return (config written)
    Utils->>Utils: Build command with --config flag
    Utils->>Visualizer: Execute ./visualize --config temp_config_path
    Visualizer->>Visualizer: Parse config_file (or use default if NULL)
    Visualizer->>Visualizer: Load environment configuration from INI
    Visualizer->>Visualizer: Generate visualization with training env settings
    Visualizer-->>Utils: Return rendered video
    Utils->>Train: Log video to wandb
Loading

Copy link

@greptile-apps greptile-apps bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

4 files reviewed, 1 comment

Edit Code Review Agent Settings | Greptile

eugenevinitsky and others added 3 commits January 5, 2026 09:52
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants