Language

unilab.algos.torch.hora.distill¶

Functions

`build_student_actor_and_normalizer`(env, cfg, ...)
`load_distilled_checkpoint`(actor, ...)
`load_teacher_actor_weights`(actor, ...)

Classes

`HoraDistillStats`	HoraDistillStats(agent_steps: 'int' = 0, best_reward: 'float' = -inf, mean_reward: 'float' = nan, mean_episode_length: 'float' = nan)
`HoraDistillationTrainer`	Stage-2 HORA latent distillation trainer.
`HoraSACDistillActor`	Stage-2 actor wrapper for HORA-SAC teachers.
`HoraSACDistillShared`	SAC-teacher-compatible HORA stage-2 shared actor.

class unilab.algos.torch.hora.distill.HoraSACDistillShared[source]¶

Bases: Module

SAC-teacher-compatible HORA stage-2 shared actor.

Parameters:

obs_dim (int)
action_dim (int)
priv_info_dim (int)
hidden_dim (int)
priv_info_embed_dim (int)
priv_mlp_hidden_dims (list[int] | tuple[int, ...])
use_layer_norm (bool)
proprio_hist_len (int)
proprio_frame_dim (int | None)
device (device | str)

__init__(obs_dim, action_dim, *, priv_info_dim, hidden_dim=512, priv_info_embed_dim=9, priv_mlp_hidden_dims=(256, 128, 9), use_layer_norm=True, proprio_hist_len=30, proprio_frame_dim=None, device='cpu')[source]¶

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Parameters:

obs_dim (int)
action_dim (int)
priv_info_dim (int)
hidden_dim (int)
priv_info_embed_dim (int)
priv_mlp_hidden_dims (list[int] | tuple[int, ...])
use_layer_norm (bool)
proprio_hist_len (int)
proprio_frame_dim (int | None)
device (device | str)

load_teacher_actor_state_dict(actor_state)[source]¶

Parameters:: actor_state (dict[str, Tensor])
Return type:: None

encode_privileged_info(priv_info)[source]¶

Parameters:: priv_info (Tensor)
Return type:: Tensor

encode_proprio_history(proprio_hist)[source]¶

Parameters:: proprio_hist (Tensor)
Return type:: Tensor

policy_mean(obs, *, prefer_student, require_privileged_target=True)[source]¶

Parameters:

obs (TensorDict)
prefer_student (bool)
require_privileged_target (bool)

Return type:

tuple[Tensor, HoraCoreOutput]

class unilab.algos.torch.hora.distill.HoraSACDistillActor[source]¶

Bases: Module

Stage-2 actor wrapper for HORA-SAC teachers.

Parameters:: shared (HoraSACDistillShared)

is_recurrent: bool = False¶

__init__(shared)[source]¶

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Parameters:: shared (HoraSACDistillShared)

forward(obs, masks=None, hidden_state=None, stochastic_output=False)[source]¶

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:

obs (TensorDict)
masks (Tensor | None)
stochastic_output (bool)

Return type:

Tensor

load_sac_teacher_actor_state_dict(actor_state)[source]¶

Parameters:: actor_state (dict[str, Tensor])
Return type:: None

class unilab.algos.torch.hora.distill.HoraDistillStats[source]¶

Bases: object

HoraDistillStats(agent_steps: ‘int’ = 0, best_reward: ‘float’ = -inf, mean_reward: ‘float’ = nan, mean_episode_length: ‘float’ = nan)

Parameters:

agent_steps (int)
best_reward (float)
mean_reward (float)
mean_episode_length (float)

agent_steps: int = 0¶

best_reward: float = -inf¶

mean_reward: float = nan¶

mean_episode_length: float = nan¶

__init__(agent_steps=0, best_reward=-inf, mean_reward=nan, mean_episode_length=nan)¶

Parameters:

agent_steps (int)
best_reward (float)
mean_reward (float)
mean_episode_length (float)

unilab.algos.torch.hora.distill.build_student_actor_and_normalizer(env, cfg, *, device)[source]¶

Parameters:

cfg (DictConfig)
device (device)

Return type:

tuple[HoraActorModel | HoraSACDistillActor, EmpiricalNormalization]

unilab.algos.torch.hora.distill.load_teacher_actor_weights(actor, teacher_checkpoint, *, teacher_algo_family, device)[source]¶

Parameters:

actor (Module)
teacher_checkpoint (str | Path)
teacher_algo_family (str)
device (device)

Return type:

None

unilab.algos.torch.hora.distill.load_distilled_checkpoint(actor, hist_normalizer, checkpoint_path, *, device)[source]¶

Parameters:

actor (Module)
hist_normalizer (EmpiricalNormalization)
checkpoint_path (str | Path)
device (device)

Return type:

dict[str, Any]

class unilab.algos.torch.hora.distill.HoraDistillationTrainer[source]¶

Bases: object

Stage-2 HORA latent distillation trainer.

Parameters:

cfg (DictConfig)
device (str)
log_dir (str | Path)
teacher_checkpoint (str | Path)
teacher_algo_family (str)
teacher_metadata (dict[str, Any] | None)
distill_runtime_cfg (DictConfig)

__init__(env, cfg, *, device, log_dir, teacher_checkpoint, teacher_algo_family, teacher_metadata=None, distill_runtime_cfg, logger)[source]¶

Parameters:

cfg (DictConfig)
device (str)
log_dir (str | Path)
teacher_checkpoint (str | Path)
teacher_algo_family (str)
teacher_metadata (dict[str, Any] | None)
distill_runtime_cfg (DictConfig)

train()[source]¶

Return type:: None

save(path)[source]¶

Parameters:: path (str | Path)
Return type:: None