Policy context loading for post-training compression.
load_float_policy_context(checkpoint_path, checkpoint_name)
Load a policy context from a float checkpoint.
Source code in src/versatil/post_training_compression/policy_loading.py
| def load_float_policy_context(
checkpoint_path: str,
checkpoint_name: str,
) -> PolicyContext:
"""Load a policy context from a float checkpoint."""
logging.info(f"Loading policy from {checkpoint_path}")
checkpoint_loader = FloatCheckpointLoader(
device=torch.device("cpu"),
checkpoint_path=checkpoint_path,
checkpoint_name=checkpoint_name,
)
return PolicyContext(
policy=checkpoint_loader.policy,
config=checkpoint_loader.config,
tokenizer=checkpoint_loader.tokenizer,
observation_space=checkpoint_loader.observation_space,
observation_horizon=checkpoint_loader.observation_horizon,
checkpoint_path=checkpoint_path,
checkpoint_name=checkpoint_name,
)
|
load_qat_policy_context(checkpoint_path, checkpoint_name, quantization)
Load a policy context from a QAT checkpoint.
Source code in src/versatil/post_training_compression/policy_loading.py
| def load_qat_policy_context(
checkpoint_path: str,
checkpoint_name: str,
quantization: BaseQuantizationWorkflow,
) -> PolicyContext:
"""Load a policy context from a QAT checkpoint."""
checkpoint_loader = _QATCheckpointLoader(
device=torch.device("cpu"),
checkpoint_path=checkpoint_path,
checkpoint_name=checkpoint_name,
quantization=quantization,
)
return PolicyContext(
policy=checkpoint_loader.policy,
config=checkpoint_loader.config,
tokenizer=checkpoint_loader.tokenizer,
observation_space=checkpoint_loader.observation_space,
observation_horizon=checkpoint_loader.observation_horizon,
checkpoint_path=checkpoint_path,
checkpoint_name=checkpoint_name,
)
|