feat(envs): add customizable state logging hook in Environment class

[srthkdev]

Description

• Introduce state_logging_hook parameter to Environment constructor
• Store state_logging_hook as an instance attribute
• Add internal _log_state method to handle state logging safely
• Integrate state logging after rollout completions in rollout methods
• Log warnings if the state logging hook raises exceptions during execution
• Enable users to inject custom state logging functionality during environment runs

Type of Change

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Documentation update
• Test improvement

Testing

• All existing tests pass
• New tests have been added to cover the changes
• Tests have been run locally with uv run pytest

Checklist

• My code follows the style guidelines of this project
• I have performed a self-review of my own code
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• Any dependent changes have been merged and published

Summary of Changes

I've implemented a custom state logging hook feature for the Verifiers library as requested in issue #330. This enhancement allows users to add custom logging functionality to track state information during environment rollouts without modifying the core library code.

Key Implementation Details

1. Environment Class Enhancement:

   • Added state_logging_hook parameter to the Environment.init method
   • Implemented _log_state() method for safe execution of the logging hook
   • Integrated state logging in both a_generate() and run_rollouts() methods

2. Safety Features:

   • Exception handling to prevent crashes if the hook fails
   • Backward compatibility (hook is optional and defaults to None)
   • Proper type checking and validation

Usage Example

def custom_state_logger(state: State) -> None:
    """Custom state logging function."""
    log_data = {
        "responses_count": len(state.get("responses", [])),
        "prompt_length": len(state.get("prompt", [])),
        "custom_fields": {k: v for k, v in state.items() if k not in ["responses", "prompt", "completion"]}
    }
    # Write to file, send to external service, etc.
    with open("state_log.jsonl", "a") as f:
        f.write(json.dumps(log_data) + "\n")

# Create environment with custom state logging hook
env = ExampleEnv(
    dataset=dataset,
    state_logging_hook=custom_state_logger
)

This implementation provides a clean, flexible, and safe way for users to add custom state logging functionality to their environments. All tests pass successfully, ensuring the implementation is robust and reliable.