Buckets:
| """Auditable per-update GRPO rollout and policy records.""" | |
| from __future__ import annotations | |
| import hashlib | |
| from pathlib import Path | |
| import torch | |
| from shared.artifacts import write_json | |
| from shared.supervised import pack_visual_codes | |
| def _span_rows(spans) -> list[list[int]]: | |
| return [[span.start, span.end] for span in spans] | |
| def _unique_attempt_root(run_dir: Path, attempted_update: int) -> Path: | |
| parent = run_dir / "rollouts" | |
| parent.mkdir(parents=True, exist_ok=True) | |
| base = f"attempt-{attempted_update:06d}" | |
| for retry in range(10_000): | |
| suffix = "" if retry == 0 else f"-retry-{retry:04d}" | |
| path = parent / f"{base}{suffix}" | |
| try: | |
| path.mkdir() | |
| return path | |
| except FileExistsError: | |
| continue | |
| raise RuntimeError(f"too many preserved RL retry artifacts for {base}") | |
| def _visual_diagnostics(trajectory) -> dict: | |
| observations = [] | |
| generated = [] | |
| generated_values = [] | |
| unavailable_grids = 0 | |
| for sequence, codes, grid in zip( | |
| trajectory.visual_sequence_spans, | |
| trajectory.visual_code_spans, | |
| trajectory.image_grid_thw, | |
| strict=True, | |
| ): | |
| value = trajectory.sampled_visual_codes[codes.start : codes.end] | |
| is_generated = sequence.start >= trajectory.prefix_length | |
| if is_generated: | |
| generated_values.append(value) | |
| height, width = int(grid[1]) // 2, int(grid[2]) // 2 | |
| if ( | |
| int(grid[0]) != 1 | |
| or int(grid[1]) % 2 | |
| or int(grid[2]) % 2 | |
| or height * width != value.shape[0] | |
| ): | |
| unavailable_grids += 1 | |
| continue | |
| target = generated if is_generated else observations | |
| target.append((value, height, width)) | |
| def resized_agreement(source, target) -> float: | |
| source_codes, source_h, source_w = source | |
| target_codes, target_h, target_w = target | |
| rows = torch.div( | |
| torch.arange(target_h) * source_h, | |
| target_h, | |
| rounding_mode="floor", | |
| ) | |
| columns = torch.div( | |
| torch.arange(target_w) * source_w, | |
| target_w, | |
| rounding_mode="floor", | |
| ) | |
| resized = source_codes.reshape( | |
| source_h, source_w, *source_codes.shape[1:] | |
| )[rows[:, None], columns[None, :]] | |
| return float( | |
| (resized.reshape_as(target_codes) == target_codes).float().mean() | |
| ) | |
| agreements = [ | |
| resized_agreement(source, target) | |
| for source in observations | |
| for target in generated | |
| ] | |
| generated_codes = ( | |
| torch.cat(generated_values) | |
| if generated_values | |
| else torch.empty(0) | |
| ) | |
| one_fraction = ( | |
| float(generated_codes.float().mean()) if generated_codes.numel() else None | |
| ) | |
| entropy_proxy = None | |
| if generated_codes.numel() and generated_codes.shape[0] > 1: | |
| probabilities = ( | |
| generated_codes.float().flatten(1).mean(0).clamp(1e-6, 1 - 1e-6) | |
| ) | |
| entropy_proxy = float( | |
| -( | |
| probabilities * probabilities.log() | |
| + (1 - probabilities) * (1 - probabilities).log() | |
| ).mean() | |
| ) | |
| offline = { | |
| "online": False, | |
| "owner": "docs/visual_quality.md", | |
| "reason": "GRPO v0 does not load an RGB decoder or image scorer", | |
| } | |
| return { | |
| "source_copy_nearest_bit_agreement": max(agreements) if agreements else None, | |
| "source_copy_comparisons": len(agreements), | |
| "source_copy_unavailable_grids": unavailable_grids, | |
| "generated_bit_one_fraction": one_fraction, | |
| "bsq_marginal_entropy_proxy": entropy_proxy, | |
| "bsq_steganography": { | |
| "online": True, | |
| "proxy": "marginal_bit_balance_and_entropy_only", | |
| "limitation": "does not detect semantic or answer-coded hidden channels", | |
| }, | |
| "blank_gray_noise_controls": offline, | |
| "ocr_text_overlay": offline, | |
| "reward_vs_shadow_metric_divergence": offline, | |
| "direct_bsq_rgb_roundtrip": offline, | |
| } | |
| def _ig_diagnostics(scored, config) -> dict: | |
| records = list(scored.ig_reward.event_records) | |
| result = {} | |
| for direction_name, direction in ( | |
| ("text_to_visual", config.reward.information_gain.text_to_visual), | |
| ("visual_to_text", config.reward.information_gain.visual_to_text), | |
| ): | |
| direction_records = [ | |
| row for row in records if row["direction"] == direction_name | |
| ] | |
| for criterion_name, component in ( | |
| ("necessity", direction.necessity), | |
| ("sufficiency", direction.sufficiency), | |
| ): | |
| eligible = [ | |
| row | |
| for row in direction_records | |
| if criterion_name == "necessity" | |
| or row["has_prior_generated_history"] | |
| ] | |
| gamma_key = f"{criterion_name}_gamma" | |
| values = [row[gamma_key] for row in eligible if gamma_key in row] | |
| blind_view = "10" if criterion_name == "necessity" else "00" | |
| full_view = "11" if criterion_name == "necessity" else "01" | |
| key = f"{direction_name}/{criterion_name}" | |
| result[key] = { | |
| "enabled": component.enabled, | |
| "eligible_events": len(eligible), | |
| "scored_events": len(values), | |
| "coverage": len(values) / len(eligible) if eligible else 0.0, | |
| "gamma_mean": sum(values) / len(values) if values else 0.0, | |
| "zero_saturation_fraction": ( | |
| sum(value == 0.0 for value in values) / len(values) | |
| if values | |
| else 0.0 | |
| ), | |
| "one_saturation_fraction": ( | |
| sum(value == 1.0 for value in values) / len(values) | |
| if values | |
| else 0.0 | |
| ), | |
| "full_score_mean": ( | |
| sum(row["scores"][full_view] for row in eligible) / len(eligible) | |
| if eligible | |
| else None | |
| ), | |
| "blind_score_mean": ( | |
| sum(row["scores"][blind_view] for row in eligible) / len(eligible) | |
| if eligible | |
| else None | |
| ), | |
| "blind_nll_mean": ( | |
| sum(row["nll"][blind_view] for row in eligible) / len(eligible) | |
| if eligible | |
| else None | |
| ), | |
| } | |
| return result | |
| def save_update_artifacts( | |
| *, | |
| run_dir: Path, | |
| attempted_update: int, | |
| groups, | |
| policy_records, | |
| config, | |
| verifier_descriptor: dict, | |
| sampling_descriptor: dict, | |
| self_scoring_descriptor: dict, | |
| ) -> Path: | |
| root = _unique_attempt_root(run_dir, attempted_update) | |
| record_index = 0 | |
| summary_groups = [] | |
| for prompt_index, group in enumerate(groups): | |
| sample_index = group.trajectories[0].trajectory.sample_index | |
| sample_id = group.trajectories[0].trajectory.sample_id | |
| sample_tag = hashlib.sha256(sample_id.encode()).hexdigest()[:12] | |
| prompt_root = root / f"prompt-{prompt_index:04d}-{sample_index:08d}-{sample_tag}" | |
| prompt_root.mkdir() | |
| summary_groups.append( | |
| { | |
| "prompt_index": prompt_index, | |
| "sample_index": sample_index, | |
| "sample_id": sample_id, | |
| "reward_mean": group.reward_mean, | |
| "reward_std": group.reward_std, | |
| "zero_advantage_reason": group.zero_advantage_reason, | |
| } | |
| ) | |
| for group_index, scored in enumerate(group.trajectories): | |
| trajectory = scored.trajectory | |
| policy = policy_records[record_index] | |
| record_index += 1 | |
| rollout_root = prompt_root / f"rollout-{group_index:02d}" | |
| rollout_root.mkdir() | |
| packed, shape = pack_visual_codes(trajectory.sampled_visual_codes) | |
| torch.save( | |
| { | |
| "schema": trajectory.schema, | |
| "input_ids": trajectory.input_ids, | |
| "prefix_length": trajectory.prefix_length, | |
| "sampled_visual_codes_packed": packed, | |
| "sampled_visual_codes_shape": shape, | |
| "image_grid_thw": trajectory.image_grid_thw, | |
| "token_owners": trajectory.token_owners, | |
| "visual_sequence_spans": _span_rows( | |
| trajectory.visual_sequence_spans | |
| ), | |
| "visual_code_spans": _span_rows(trajectory.visual_code_spans), | |
| "text_action_positions": trajectory.text_action_positions, | |
| "visual_action_positions": trajectory.visual_action_positions, | |
| "text_old_log_probs": trajectory.text_old_log_probs, | |
| "visual_old_factor_log_probs": ( | |
| trajectory.visual_old_factor_log_probs | |
| ), | |
| }, | |
| rollout_root / "trajectory.pt", | |
| ) | |
| torch.save( | |
| { | |
| "log_probability_semantics": ( | |
| "sampling_old_value_aligned_full_prefix_gradient_" | |
| "surrogate_v1" | |
| ), | |
| "text_aligned_surrogate_log_probs": ( | |
| policy.text_current_log_probs | |
| ), | |
| "text_old_log_probs": policy.text_old_log_probs, | |
| "text_reference_log_probs": policy.text_reference_log_probs, | |
| "text_full_prefix_replay_log_probs": ( | |
| policy.text_replay_log_probs | |
| ), | |
| "text_aligned_surrogate_old_log_ratio": ( | |
| policy.text_current_log_probs - policy.text_old_log_probs | |
| ), | |
| "text_full_prefix_replay_reference_log_ratio": ( | |
| policy.text_replay_log_probs | |
| - policy.text_reference_log_probs | |
| ), | |
| "visual_aligned_surrogate_factor_log_probs": ( | |
| policy.visual_current_factor_log_probs | |
| ), | |
| "visual_old_factor_log_probs": ( | |
| policy.visual_old_factor_log_probs | |
| ), | |
| "visual_reference_factor_log_probs": ( | |
| policy.visual_reference_factor_log_probs | |
| ), | |
| "visual_full_prefix_replay_factor_log_probs": ( | |
| policy.visual_replay_factor_log_probs | |
| ), | |
| "visual_aligned_surrogate_old_log_ratio": ( | |
| policy.visual_current_factor_log_probs | |
| - policy.visual_old_factor_log_probs | |
| ), | |
| "visual_full_prefix_replay_reference_log_ratio": ( | |
| policy.visual_replay_factor_log_probs | |
| - policy.visual_reference_factor_log_probs | |
| ), | |
| }, | |
| rollout_root / "policy.pt", | |
| ) | |
| ig_components = { | |
| f"{direction}/{criterion}": value | |
| for (direction, criterion), value in scored.ig_reward.components.items() | |
| } | |
| result = { | |
| "sample_id": trajectory.sample_id, | |
| "sample_index": trajectory.sample_index, | |
| "group_index": trajectory.group_index, | |
| "stop_reason": trajectory.stop_reason, | |
| "metadata": dict(trajectory.metadata), | |
| "protocol": { | |
| "valid": scored.protocol.valid, | |
| "failure_reason": scored.protocol.failure_reason, | |
| "final_answer_text": scored.protocol.final_answer_text, | |
| "blocks": len(scored.protocol.blocks), | |
| "transitions": len(scored.protocol.transitions), | |
| }, | |
| "verifier": verifier_descriptor, | |
| "sampling": sampling_descriptor, | |
| "verifier_correct": scored.verifier_correct, | |
| "reward": { | |
| "format": scored.reward.format, | |
| "accuracy": scored.reward.accuracy, | |
| "image": scored.reward.image, | |
| "information_gain": scored.reward.ig, | |
| "cost": scored.reward.cost, | |
| "total": scored.reward.total, | |
| "ig_components": ig_components, | |
| "normalized_cost": scored.normalized_cost, | |
| }, | |
| "four_view_self_scoring": { | |
| "enabled": config.reward.information_gain.enabled, | |
| "scorer": ( | |
| "pre_update_behavior" | |
| if config.reward.information_gain.enabled | |
| else "disabled_by_zero_outer_weight" | |
| ), | |
| "events": scored.ig_reward.event_records, | |
| "descriptor": self_scoring_descriptor, | |
| "component_diagnostics": _ig_diagnostics(scored, config), | |
| }, | |
| "group": { | |
| "mean": group.reward_mean, | |
| "std": group.reward_std, | |
| "advantage": float(group.advantages[group_index]), | |
| "zero_advantage_reason": group.zero_advantage_reason, | |
| }, | |
| "policy": { | |
| "log_probability_semantics": ( | |
| "sampling_old_value_aligned_full_prefix_gradient_" | |
| "surrogate_v1" | |
| ), | |
| "behavior_log_prob_max_abs_error": ( | |
| policy.behavior_log_prob_max_abs_error | |
| ), | |
| "text_behavior_log_prob_max_abs_error": ( | |
| policy.text_behavior_log_prob_max_abs_error | |
| ), | |
| "visual_behavior_log_prob_max_abs_error": ( | |
| policy.visual_behavior_log_prob_max_abs_error | |
| ), | |
| "text_behavior_log_prob_mean_abs_error": ( | |
| policy.text_behavior_log_prob_mean_abs_error | |
| ), | |
| "text_behavior_log_prob_p99_abs_error": ( | |
| policy.text_behavior_log_prob_p99_abs_error | |
| ), | |
| "visual_behavior_log_prob_mean_abs_error": ( | |
| policy.visual_behavior_log_prob_mean_abs_error | |
| ), | |
| "visual_behavior_log_prob_p99_abs_error": ( | |
| policy.visual_behavior_log_prob_p99_abs_error | |
| ), | |
| "visual_behavior_log_prob_step_signed_mean_abs_p99": ( | |
| policy.visual_behavior_log_prob_step_signed_mean_abs_p99 | |
| ), | |
| "text_actions": policy.text_current_log_probs.numel(), | |
| "visual_steps": policy.visual_current_factor_log_probs.shape[0], | |
| "text_ratio_mean": policy.text_ratio_mean, | |
| "visual_factor_ratio_mean": policy.visual_factor_ratio_mean, | |
| "text_clip_fraction": policy.text_clip_fraction, | |
| "visual_factor_clip_fraction": ( | |
| policy.visual_factor_clip_fraction | |
| ), | |
| "reference_jsd": policy.reference_jsd, | |
| "autonomous_ar_steps": policy.autonomous_ar_steps, | |
| }, | |
| "active_components": { | |
| "format": config.reward.format.weight > 0, | |
| "accuracy": config.reward.accuracy.weight > 0, | |
| "information_gain": config.reward.information_gain.enabled, | |
| "cost": config.reward.cost.enabled, | |
| }, | |
| "disabled_components": { | |
| "image": group.image_reward_descriptor, | |
| }, | |
| "visual_diagnostics": _visual_diagnostics(trajectory), | |
| } | |
| write_json(rollout_root / "result.json", result) | |
| if record_index != len(policy_records): | |
| raise ValueError("policy records do not align with sampled trajectories") | |
| write_json( | |
| root / "update.json", | |
| { | |
| "attempted_update": attempted_update, | |
| "groups": summary_groups, | |
| "sampling": { | |
| "group_size": config.rollout.group_size, | |
| "text_temperature": config.rollout.text_temperature, | |
| "visual_temperature": config.rollout.visual_temperature, | |
| "top_p": config.rollout.top_p, | |
| "top_k": config.rollout.top_k, | |
| "cfg": config.rollout.cfg_scale, | |
| "fingerprint": sampling_descriptor["fingerprint"], | |
| }, | |
| "self_scoring": self_scoring_descriptor, | |
| "image_reward": {"backend": "disabled", "value": 0.0}, | |
| }, | |
| ) | |
| return root | |
| __all__ = ["save_update_artifacts"] | |
Xet Storage Details
- Size:
- 17.2 kB
- Xet hash:
- 450ae102bbd4f027e2e3c50ce480127aeb32de245f36091c6835b82688ef0f47
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.