float_runtime

float_runtime

Floating-point policy inference runtime.

FloatPolicyRuntime

FloatPolicyRuntime(device, checkpoint_path, checkpoint_name=value, seed=42, compile_model=True)

Bases: PolicyRuntime

Inference runtime for floating-point policy checkpoints.

Initialize the float policy runtime.

Parameters:

Name	Type	Description	Default
device	device	Device to load the model onto.	required
checkpoint_path	str	Path to the checkpoint directory.	required
checkpoint_name	str	Name of the checkpoint file.	value
seed	int	Random seed for reproducibility.	42
compile_model	bool	Whether to compile the policy with torch.compile for optimized inference.	True

Source code in src/versatil/inference/policy_runtime/float_runtime.py

def __init__(
    self,
    device: torch.device,
    checkpoint_path: str,
    checkpoint_name: str = CheckpointFilename.DEFAULT_CHECKPOINT.value,
    seed: int = 42,
    compile_model: bool = True,
) -> None:
    """Initialize the float policy runtime.

    Args:
        device: Device to load the model onto.
        checkpoint_path: Path to the checkpoint directory.
        checkpoint_name: Name of the checkpoint file.
        seed: Random seed for reproducibility.
        compile_model: Whether to compile the policy with
            torch.compile for optimized inference.
    """
    self._compile_model = compile_model
    self._set_seed(seed)
    checkpoint_loader = FloatCheckpointLoader(
        device=device,
        checkpoint_path=checkpoint_path,
        checkpoint_name=checkpoint_name,
    )
    super().__init__(
        checkpoint_loader=checkpoint_loader,
        client_identifier=str(Path(checkpoint_path) / Path(checkpoint_name).stem),
    )
    self._prepare_inference_model()

run_inference

run_inference(obs_dict)

Run policy inference with autocast and no_grad.

Parameters:

Name	Type	Description	Default
obs_dict	dict[str, Tensor]	Observation dictionary for the policy.	required

Returns:

Type	Description
dict[str, Tensor]	Action dictionary from policy.predict_action.

Source code in src/versatil/inference/policy_runtime/float_runtime.py

def run_inference(
    self, obs_dict: dict[str, torch.Tensor]
) -> dict[str, torch.Tensor]:
    """Run policy inference with autocast and no_grad.

    Args:
        obs_dict: Observation dictionary for the policy.

    Returns:
        Action dictionary from policy.predict_action.
    """
    with (
        torch.autocast(
            device_type=self.device.type,
            dtype=PrecisionType(self._precision).get_model_dtype(),
        ),
        torch.no_grad(),
    ):
        return self._policy.predict_action(obs_dict=obs_dict)