# conversations.py """Conversation builders for RynnValueLangProcessor. Each builder turns ``(instruction, num_frames, ...)`` into the chat-template content list and tells the processor how many ```` / ```` tokens to expect for each conversation type. Three conversation types are supported: * ``progress`` — predict remaining time to task completion via ````; optionally appends the success block. * ``relative`` — predict remaining time to a target state via per-frame ````. * ``progress_relative`` — progress prompt followed by a trailing target image and a single ```` prediction (multi-task supervision). The processor stays agnostic to the prompt layout — it only reads ``progress_value_count`` to align target tensors. """ from abc import ABC, abstractmethod from typing import List, Optional from PIL.Image import Image def _meta_block( robot_description: Optional[str], camera_description: Optional[str], ) -> List[dict]: sentences = [] if robot_description is not None: sentences.append(f"The agent is {robot_description}.") if camera_description is not None: sentences.append(f"The observation is captured from {camera_description}.") return [{"type": "text", "text": " ".join(sentences)}] def _analysis_block( description: Optional[str] = None, match_answer: Optional[str] = None, success_answer: Optional[str] = None, ) -> List[dict]: lines = [] if description: lines.append(f"- Video Description: {description}") if match_answer is not None: lines.append(f"- Match: {match_answer}") if success_answer is not None: lines.append(f"- Success: {success_answer}") if not lines: return [] return [ {"type": "text", "text": "Analysis: \n"}, {"type": "text", "text": "\n".join(lines)}, ] class ConversationBuilder(ABC): """Strategy object that owns the prompt layout for one processor. Subclasses implement ``build_progress`` and ``build_relative`` and expose ``progress_value_count(num_frames)`` so the processor knows how many ```` tokens each progress sample contains. """ name: str = "base" PROGRESS_QUESTION = ( "Estimate the minimum remaining time in seconds until the agent completes the task." ) RELATIVE_QUESTION = ( "Estimate the minimum remaining time in seconds until the agent reaches the target state." ) def __init__( self, value_token: str, relative_value_token: Optional[str], use_meta: bool = False, value_token_repeat: int = 1, relative_value_token_repeat: int = 1, ): self.value_token = value_token self.relative_value_token = relative_value_token self.use_meta = use_meta if value_token_repeat < 1: raise ValueError( f"value_token_repeat must be >= 1, got {value_token_repeat}" ) self.value_token_repeat = int(value_token_repeat) if relative_value_token_repeat < 1: raise ValueError( f"relative_value_token_repeat must be >= 1, got {relative_value_token_repeat}" ) self.relative_value_token_repeat = int(relative_value_token_repeat) def _maybe_meta( self, robot_description: Optional[str], camera_description: Optional[str], ) -> List[dict]: if not self.use_meta: return [] if robot_description is None and camera_description is None: raise ValueError( "use_meta=True requires at least one of `robot_description` " "or `camera_description` to be provided." ) return _meta_block(robot_description, camera_description) def _instruction_block(self, instruction: str) -> List[dict]: return [ {"type": "text", "text": f"The agent is performing the following task: {instruction}."}, ] def _append_value_tokens(self, content: List[dict]) -> None: """Append ``value_token_repeat`` copies of ````. Repeated tokens supply extra positions per prediction slot; the model mean-pools their hidden states before feeding the value head. """ for _ in range(self.value_token_repeat): content.append({"type": "text", "text": f"{self.value_token}"}) def _append_relative_value_tokens(self, content: List[dict]) -> None: """Append ``relative_value_token_repeat`` copies of ````. Mirrors :meth:`_append_value_tokens` for the relative head; the model mean-pools the R repeated hidden states into one prediction per slot. """ for _ in range(self.relative_value_token_repeat): content.append({"type": "text", "text": f"{self.relative_value_token}"}) @abstractmethod def progress_value_count(self, num_frames: int) -> int: """Number of prediction slots (targets) emitted by the progress prompt. Each slot may be represented by ``value_token_repeat`` consecutive ```` tokens in the prompt; the model pools them back down to one prediction per slot, so this returns the slot/target count, not the raw token count. """ @abstractmethod def build_progress(self, **kwargs) -> List[dict]: """Return the chat-template content list for a progress sample.""" def build_relative( self, instruction: str, num_frames: int, target_image: Image, robot_description: Optional[str], camera_description: Optional[str], ) -> List[dict]: if self.relative_value_token is None: raise ValueError( "relative conversation requires a registered token " "(enable relative_value_head_config). Relative mode predicts only the " "relative value (time-to-target) and never the step ." ) if target_image is None: raise ValueError("relative conversation requires `target_image`.") content: List[dict] = self._instruction_block(instruction) content.extend(self._maybe_meta(robot_description, camera_description)) content.extend([ {"type": "image"}, {"type": "text", "text": "The image above shows the target state the agent needs to reach."}, ]) content.append({"type": "text", "text": f"Question: {self.RELATIVE_QUESTION}"}) for _ in range(num_frames): content.append({"type": "image"}) self._append_relative_value_tokens(content) return content def build_progress_relative( self, instruction: str, num_frames: int, target_image: Image, robot_description: Optional[str], camera_description: Optional[str], ) -> List[dict]: """Progress prompt followed by a trailing relative-value prediction. Reuses :meth:`build_progress` for the prefix and appends a target image plus a single ```` token, supervised against the time-to-target. Independent of the concrete progress layout, so subclasses inherit this default. """ if self.relative_value_token is None: raise ValueError( "progress_relative conversation requires a registered " " token (enable relative_value_head_config)." ) if target_image is None: raise ValueError("progress_relative conversation requires `target_image`.") content = self.build_progress( instruction=instruction, num_frames=num_frames, robot_description=robot_description, camera_description=camera_description, ) prediction_tokens = " ".join( [self.relative_value_token] * self.relative_value_token_repeat ) content.extend([ {"type": "image"}, {"type": "text", "text": "The image above shows the target state the agent needs to reach."}, {"type": "text", "text": f"Question: {self.RELATIVE_QUESTION}"}, {"type": "text", "text": f"Prediction: {prediction_tokens}"}, ]) return content class InterleavedConversationBuilder(ConversationBuilder): """Legacy layout: instruction + meta + question, then per-frame [img]. Each frame carries its own ```` token, so progress targets are per-frame remaining times. """ name = "interleaved" def progress_value_count(self, num_frames: int) -> int: return num_frames def build_progress( self, instruction: str, num_frames: int, robot_description: Optional[str], camera_description: Optional[str], ) -> List[dict]: content: List[dict] = self._instruction_block(instruction) content.extend(self._maybe_meta(robot_description, camera_description)) content.append({"type": "text", "text": f"Question: {self.PROGRESS_QUESTION}"}) for _ in range(num_frames): content.append({"type": "image"}) self._append_value_tokens(content) return content class InterleavedLangConversationBuilder(ConversationBuilder): """Interleaved layout with language + relative distance + value prediction. Per-frame layout: ``[img] {instruction} `` - ``{instruction}`` is the language prediction target (causal LM loss). - ```` predicts the signed time delta from previous frame. - ```` predicts the remaining time to task completion. Layout: [meta] + questions + per-frame([img] instruction ) """ name = "interleaved_lang" LANG_QUESTION = "For each frame, what task is the agent performing?" RELATIVE_DISTANCE_QUESTION = "For each frame, what is the time elapsed since the previous frame?" def progress_value_count(self, num_frames: int) -> int: return num_frames def build_progress( self, instructions: List[str], num_frames: int, robot_description: Optional[str], camera_description: Optional[str], ) -> List[dict]: if len(instructions) != num_frames: raise ValueError( f"instructions length ({len(instructions)}) must equal " f"num_frames ({num_frames})." ) if self.relative_value_token is None: raise ValueError( "interleaved_lang conversation requires a registered " " token (enable relative_value_head_config)." ) content: List[dict] = [] content.extend(self._maybe_meta(robot_description, camera_description)) content.append({"type": "text", "text": f"Question: {self.PROGRESS_QUESTION}"}) content.append({"type": "text", "text": f"Question: {self.LANG_QUESTION}"}) content.append({"type": "text", "text": f"Question: {self.RELATIVE_DISTANCE_QUESTION}"}) for i in range(num_frames): content.append({"type": "image"}) content.append({"type": "text", "text": f"Task: {instructions[i]};"}) self._append_relative_value_tokens(content) self._append_value_tokens(content) return content class InterleavedHistoryConversationBuilder(ConversationBuilder): """History layout with per-frame value + relative-value tokens + trailing analysis. Layout:: [meta] The agent is performing the following task: {instruction}. Question: Question: # per-frame, i = 0..N-1: [img] (if i > 0) ×R ×R Analyze this trajectory in terms of video description, match, and success. Analysis: Video Description: {description} Match: Yes/No Success: Yes/No """ name = "interleaved_history" RELATIVE_DISTANCE_QUESTION = ( "For each frame after the first, what is the time delta from the previous frame?" ) @staticmethod def _build_analysis_prompt(has_description: bool) -> str: lines = ["Analyze this trajectory. Provide:"] if has_description: lines.append("- Video Description: a brief description of what the agent is doing in the video.") lines.append("- Match: whether the video matches the stated task (Yes/No).") lines.append("- Success: whether the agent has completed the task (Yes/No).") return "\n".join(lines) def progress_value_count(self, num_frames: int) -> int: return num_frames def build_progress( self, instruction: str, description: Optional[str], num_frames: int, robot_description: Optional[str], camera_description: Optional[str], match_answer: Optional[str] = None, success_answer: Optional[str] = None, is_inference: bool = False, ) -> List[dict]: use_relative = self.relative_value_token is not None content: List[dict] = [] content.extend(self._maybe_meta(robot_description, camera_description)) content.extend(self._instruction_block(instruction)) if use_relative: content.append({"type": "text", "text": f"Question: {self.RELATIVE_DISTANCE_QUESTION}"}) content.append({"type": "text", "text": f"Question: {self.PROGRESS_QUESTION}"}) for i in range(num_frames): content.append({"type": "image"}) if i > 0 and use_relative: self._append_relative_value_tokens(content) self._append_value_tokens(content) # Inference: always include all questions so the model sees full context. # Training: only include description question when description data exists. prompt_has_description = True if is_inference else bool(description) analysis_prompt = self._build_analysis_prompt(has_description=prompt_has_description) analysis = _analysis_block( description=description, match_answer=match_answer, success_answer=success_answer, ) if not analysis and is_inference: # Inference: emit the analysis prompt questions together with # the "Analysis: \n" marker so generation has full context. content.append({"type": "text", "text": analysis_prompt}) analysis = [{"type": "text", "text": "Analysis: \n"}] else: content.append({"type": "text", "text": analysis_prompt}) content.extend(analysis) return content class InterleavedHistoryNormalizedConversationBuilder(InterleavedHistoryConversationBuilder): """Same layout as :class:`InterleavedHistoryConversationBuilder`, but the prompt asks for a normalized [0, 1] completion progress instead of the raw remaining time in seconds. Only the question strings differ; the frame / token layout, value counts and analysis block are inherited unchanged, so the model still emits one ```` per frame and ```` between frames — the targets are just normalized upstream in the processor. """ name = "interleaved_history_normalized" PROGRESS_QUESTION = ( "Estimate the task completion progress as a value between 0 and 1, where " "0 means the task has just started and 1 means the task is complete." ) RELATIVE_DISTANCE_QUESTION = ( "For each frame after the first, what fraction of the total task time has " "elapsed since the previous frame?" ) class VideoPrefixConversationBuilder(ConversationBuilder): """New layout: meta + "this is a video" preface + frames (no value tokens), then task description + question + a single trailing ```` token, optionally followed by the success block. Progress targets become a single anchor (last-frame) remaining time. """ name = "video_prefix" INTRO_TEXT = "The following is a video of the agent completing a manipulation task." def progress_value_count(self, num_frames: int) -> int: return 1 def build_progress( self, instruction: str, num_frames: int, robot_description: Optional[str], camera_description: Optional[str], ) -> List[dict]: if self.relative_value_token is None: raise ValueError( "video_prefix conversation requires a registered " " token (enable relative_value_head_config)." ) content: List[dict] = [] content.extend(self._maybe_meta(robot_description, camera_description)) content.append({"type": "text", "text": self.INTRO_TEXT}) for _ in range(num_frames): content.append({"type": "image"}) self._append_relative_value_tokens(content) content.append({"type": "text", "text": f"The task is: {instruction}."}) content.append({"type": "text", "text": f"Question: {self.PROGRESS_QUESTION}"}) prediction_tokens = " ".join([self.value_token] * self.value_token_repeat) content.append({"type": "text", "text": f"Prediction: {prediction_tokens}"}) return content _REGISTRY = { InterleavedConversationBuilder.name: InterleavedConversationBuilder, InterleavedLangConversationBuilder.name: InterleavedLangConversationBuilder, InterleavedHistoryConversationBuilder.name: InterleavedHistoryConversationBuilder, InterleavedHistoryNormalizedConversationBuilder.name: InterleavedHistoryNormalizedConversationBuilder, VideoPrefixConversationBuilder.name: VideoPrefixConversationBuilder, } def build_conversation_builder( style: str, *, value_token: str, relative_value_token: Optional[str], use_meta: bool = False, value_token_repeat: int = 1, relative_value_token_repeat: int = 1, ) -> ConversationBuilder: if style not in _REGISTRY: raise ValueError( f"Unknown conversation_style: {style!r}. " f"Expected one of {sorted(_REGISTRY)}." ) return _REGISTRY[style]( value_token=value_token, relative_value_token=relative_value_token, use_meta=use_meta, value_token_repeat=value_token_repeat, relative_value_token_repeat=relative_value_token_repeat, ) __all__ = [ "ConversationBuilder", "InterleavedConversationBuilder", "InterleavedLangConversationBuilder", "InterleavedHistoryConversationBuilder", "InterleavedHistoryNormalizedConversationBuilder", "VideoPrefixConversationBuilder", "build_conversation_builder", ]