lightrft.utils.trajectory_saver

Trajectory Saver Utility for debugging and analysis.

This module provides utilities to save experience trajectories to JSON files for debugging and analysis purposes.

class lightrft.utils.trajectory_saver.TrajectorySaver(save_dir: str, tokenizer: Any, save_images_separately: bool = True, max_image_size: int = 512, mark_high_entropy_tokens: bool = False, high_entropy_token_ratio: float = 0.2)

Bases: object

Utility class to save experience trajectories to JSON files.

Features:

• Saves experience sequences, rewards, and metadata for individual samples.

• Supports both text-only and vision-language models.

• Efficiently handles image data by saving them to a separate directory with clear linkage.

• Only saves on rank 0 to avoid duplication.

• Produces human-readable JSON output for easy debugging.

Parameters:

• save_dir (str) – Directory to save trajectory files

• tokenizer (Any) – Tokenizer for decoding sequences

• save_images_separately (bool) – If True, save images as separate files. Default to True

• max_image_size (int) – Maximum dimension for saved images (to reduce file size). Default to 512

• mark_high_entropy_tokens (bool) – If True, create token arrays with high-entropy information for HTML rendering. Default to False

• high_entropy_token_ratio (float) – Ratio of high-entropy tokens to mark (e.g., 0.2 means top 20%). Only used if mark_high_entropy_tokens is True. Default to 0.2

save_trajectories(experiences: List[Any], step: int, num_samples: int = 10, prefix: str = 'trajectories', compute_stats: bool = False) → Tuple[str | None, Dict[str, float] | None]

Save a subset of experiences to a JSON file and optionally compute statistics.

Each Experience object is a micro-batch. This function unpacks them into individual sample trajectories before saving.

Parameters:

• experiences (List[Any]) – List of Experience or ExperienceVL objects from the replay buffer

• step (int) – Current training step (used in filename)

• num_samples (int) – Target number of individual sample trajectories to save. Default to 10

• prefix (str) – Prefix for the output filename. Default to “trajectories”

• compute_stats (bool) – Whether to compute and return statistics. Default to False

Returns:

Tuple of (path to saved JSON file, statistics dict) or (None, None) if not rank 0 or no experiences

Return type:

Tuple[Optional[str], Optional[Dict[str, float]]]

lightrft.utils.trajectory_saver.create_trajectory_saver(args: Any, tokenizer: Any) → TrajectorySaver | None

Factory function to create TrajectorySaver if enabled.

Parameters:

• args (Any) – Training arguments containing save_trajectories flag and save_path

• tokenizer (Any) – Tokenizer for decoding sequences

Returns:

TrajectorySaver instance or None if not enabled

Return type:

Optional[TrajectorySaver]