gm.nn.AnchoredPolicy#

class gemma.gm.nn.AnchoredPolicy( policy: flax.linen.module.Module, anchor: flax.linen.module.Module | None = None, parent: flax.linen.module.Module | flax.core.scope.Scope | flax.linen.module._Sentinel | None = <flax.linen.module._Sentinel object>, name: str | None = None, )[source]

Bases: flax.linen.module.Module

Wrapper around a model to compute policy and anchor outputs.

This wrapper takes an input and pass it through two models: - policy: Model trained. - anchor: Frozen model. If not provided, is set to a copy of policy.

To initialize the model, use gm.ckpts.AnchoredPolicyLoader.

policy: flax.linen.module.Module

anchor: flax.linen.module.Module | None = None

name: str | None = None

parent: flax.linen.module.Module | flax.core.scope.Scope | flax.linen.module._Sentinel | None = None

scope: flax.core.scope.Scope | None = None