Feature/atari rendering

[giongto35]

Summary

This PR enhances the Atari environment's web interface with visual rendering capabilities and simplifies the example scripts. It adds a toggle to show/hide visual game rendering in the web UI (Atari-only) and converts atari_simple.py to use proper async/await patterns for better reliability.

Type of Change

• New feature
• Refactoring

Alignment Checklist

Before submitting, verify:

• I have read .claude/docs/PRINCIPLES.md and this PR aligns with our principles
• I have checked .claude/docs/INVARIANTS.md and no invariants are violated
• I have run /pre-submit-pr (or bash .claude/hooks/lint.sh and tests) and addressed all issues

RFC Status

• Not required (bug fix, docs, minor refactoring)

Changes

Web Interface Enhancements

• Added visual rendering toggle for Atari environments: The web interface now includes a "Render Visual" checkbox that allows users to toggle between viewing the raw JSON observation data and the visual game screen
• Conditional rendering: The toggle only appears for Atari environments (detected via 'atari' in metadata.name.lower())
• Improved UX: Visual rendering is enabled by default for Atari, making it easier to see gameplay in real-time

Example Script Improvements

• Converted atari_simple.py to async/await: Removed the problematic .sync() wrapper that was causing "Future attached to a different loop" errors with persistent WebSocket connections
• Added --debug flag: Allows running against a local server (http://127.0.0.1:8001) instead of Docker for faster development iteration
• Removed unnecessary files: Cleaned up atari_interactive.py and atari_example.py to reduce maintenance burden

Code Quality

• Better async handling: Uses native async def main() with asyncio.run() instead of sync wrappers
• Simplified imports: Uses direct from atari_env import AtariEnv, AtariAction instead of inline class definitions

Test Plan

Manual Testing

1. Start the Atari server:

   export ENABLE_WEB_INTERFACE=true && \
   export ATARI_GAME=breakout && \
   export PYTHONPATH=$(pwd)/envs:$(pwd)/src:$(pwd) && \
   uvicorn envs.atari_env.server.app:app --host 0.0.0.0 --port 8001

[greptile-apps]

Greptile Overview

Greptile Summary

This PR adds visual rendering capabilities to the Atari environment's web interface and improves the example scripts with proper async/await patterns.

Key Changes:

• Added screen_image field to AtariObservation for base64-encoded PNG visualization
• Server now generates base64-encoded PNG images from RGB screen data using Pillow
• Web interface includes Atari-specific toggle to switch between visual rendering and JSON observation display
• atari_simple.py converted from sync wrapper to native async/await, fixing "Future attached to a different loop" errors
• Added --debug flag to connect to local server instead of Docker

Issues Found:

• Print statement used for error handling in atari_environment.py:264 (debug code in production)
• Missing screen_image field extraction in client.py _parse_result() method - the new field won't be available to client code

No Alignment Concerns:

• Changes are isolated to Atari environment implementation
• No violations of client-server separation, API signatures, or security invariants
• Web interface changes are presentation-only and don't expose simulation controls
• Async conversion follows established patterns in the codebase

Confidence Score: 3/5

• Safe to merge after fixing the missing field extraction in client.py
• Two concrete issues found: print statement in production code (minor) and missing screen_image field in client parser (will cause the new feature to not work in client code). The async conversion is solid and follows best practices. No architectural concerns.
• envs/atari_env/server/atari_environment.py and envs/atari_env/client.py need fixes before merge

Important Files Changed

Filename	Overview
envs/atari_env/server/atari_environment.py	Added screen image encoding with print statement for error handling
examples/atari_simple.py	Converted to async/await pattern and added --debug flag for local development
envs/atari_env/client.py	Missing screen_image field extraction in _parse_result method - new feature won't work

Sequence Diagram

sequenceDiagram
    participant User
    participant AtariSimple as atari_simple.py
    participant AtariEnv as AtariEnv Client
    participant WebSocket as WebSocket Connection
    participant Server as Atari Environment Server
    participant WebUI as Web Interface

    User->>AtariSimple: python atari_simple.py [--debug]
    
    alt Debug Mode
        AtariSimple->>AtariEnv: AtariEnv(base_url="http://127.0.0.1:8001")
    else Standard Mode
        AtariSimple->>AtariEnv: await from_docker_image()
        AtariEnv->>Server: Start Docker container
    end
    
    AtariSimple->>AtariEnv: await connect()
    AtariEnv->>WebSocket: Establish WebSocket connection
    WebSocket->>Server: Connect
    
    AtariSimple->>AtariEnv: await reset()
    AtariEnv->>WebSocket: Send reset request
    WebSocket->>Server: reset()
    Server->>Server: Generate RGB screen
    Server->>Server: Encode screen to base64 PNG
    Server-->>WebSocket: AtariObservation + screen_image
    WebSocket-->>AtariEnv: Observation with screen_image
    AtariEnv-->>AtariSimple: StepResult
    
    loop For each step
        AtariSimple->>AtariEnv: await step(AtariAction)
        AtariEnv->>WebSocket: Send action
        WebSocket->>Server: step(action)
        Server->>Server: Execute action with frameskip
        Server->>Server: Generate screen + encode to PNG
        Server-->>WebSocket: Observation + screen_image
        WebSocket-->>AtariEnv: Observation
        AtariEnv-->>AtariSimple: StepResult
    end
    
    Note over WebUI: Parallel Web Interface Flow
    User->>WebUI: Open web interface in browser
    WebUI->>Server: WebSocket connection
    WebUI->>WebUI: Check if 'atari' in env name
    WebUI->>WebUI: Show "Render Visual" toggle
    
    loop On observation update
        Server-->>WebUI: state_update with screen_image
        alt Toggle checked
            WebUI->>WebUI: Display <img> with screen_image
        else Toggle unchecked
            WebUI->>WebUI: Display JSON observation
        end
    end
    
    AtariSimple->>AtariEnv: await close()
    AtariEnv->>WebSocket: Close connection
    WebSocket->>Server: Disconnect

Loading

[greptile-apps]

_{3 files reviewed, 2 comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

Additional Comments (1)

envs/atari_env/client.py
Missing screen_image field extraction from observation data. The new screen_image field added to AtariObservation in models.py isn't being parsed here

        observation = AtariObservation(
            screen=obs_data.get("screen", []),
            screen_shape=obs_data.get("screen_shape", []),
            legal_actions=obs_data.get("legal_actions", []),
            lives=obs_data.get("lives", 0),
            episode_frame_number=obs_data.get("episode_frame_number", 0),
            frame_number=obs_data.get("frame_number", 0),
            screen_image=obs_data.get("screen_image"),
            done=payload.get("done", False),
            reward=payload.get("reward"),
            metadata=obs_data.get("metadata", {}),
        )

Prompt To Fix With AI

This is a comment left during a code review.
Path: envs/atari_env/client.py
Line: 83:93

Comment:
Missing `screen_image` field extraction from observation data. The new `screen_image` field added to `AtariObservation` in models.py isn't being parsed here

```suggestion
        observation = AtariObservation(
            screen=obs_data.get("screen", []),
            screen_shape=obs_data.get("screen_shape", []),
            legal_actions=obs_data.get("legal_actions", []),
            lives=obs_data.get("lives", 0),
            episode_frame_number=obs_data.get("episode_frame_number", 0),
            frame_number=obs_data.get("frame_number", 0),
            screen_image=obs_data.get("screen_image"),
            done=payload.get("done", False),
            reward=payload.get("reward"),
            metadata=obs_data.get("metadata", {}),
        )
```

How can I resolve this? If you propose a fix, please make it concise.

[burtenshaw]

Happy to merge this if you can respond to greptile and lint.

[giongto35]

Fixed greptile review