Deploy a trained PPO (shared model)

[AndreaRossetto]

Hi,
I'm trying t odeploy a trained PPO agent, the chekpoint is shered (policy + value).
I create the shared model (I get the code during the training in isaaclab)

class SharedModel(GaussianMixin,DeterministicMixin, Model):
    def __init__(self, observation_space, action_space, device):
        Model.__init__(self, observation_space, action_space, device)
        GaussianMixin.__init__(
            self,
            clip_actions=False,
            clip_log_std=True,
            min_log_std=-20.0,
            max_log_std=2.0,
            reduction="sum",
            role="policy",
        )
        DeterministicMixin.__init__(self, clip_actions=False, role="value")

        self.net_container = nn.Sequential(
            nn.LazyLinear(out_features=512),
            nn.ELU(),
            nn.LazyLinear(out_features=256),
            nn.ELU(),
            nn.LazyLinear(out_features=128),
            nn.ELU(),
        )
        self.policy_layer = nn.LazyLinear(out_features=self.num_actions)
        self.log_std_parameter = nn.Parameter(torch.full(size=(self.num_actions,), fill_value=0.0), requires_grad=True)
        self.value_layer = nn.LazyLinear(out_features=1)

    def act(self, inputs, role):
        if role == "policy":
            return GaussianMixin.act(self, inputs, role)
        elif role == "value":
            return DeterministicMixin.act(self, inputs, role)

    def compute(self, inputs, role=""):
        if role == "policy":
            states = unflatten_tensorized_space(self.observation_space, inputs.get("states"))
            taken_actions = unflatten_tensorized_space(self.action_space, inputs.get("taken_actions"))
            net = self.net_container(states)
            self._shared_output = net
            output = self.policy_layer(net)
            return output, self.log_std_parameter, {}
        elif role == "value":
            if self._shared_output is None:
                states = unflatten_tensorized_space(self.observation_space, inputs.get("states"))
                taken_actions = unflatten_tensorized_space(self.action_space, inputs.get("taken_actions"))
                net = self.net_container(states)
                shared_output = net
            else:
                shared_output = self._shared_output
            self._shared_output = None
            output = self.value_layer(shared_output)
            return output, {}

when I try to run it I get this error

RuntimeError: Error(s) in loading state_dict for SharedModel:
        Missing key(s) in state_dict: "log_std_parameter", "net_container.0.weight", "net_container.0.bias", "net_container.2.weight", "net_container.2.bias", "net_container.4.weight", "net_container.4.bias", "policy_layer.weight", "policy_layer.bias", "value_layer.weight", "value_layer.bias". 
        Unexpected key(s) in state_dict: "model_state_dict", "optimizer_state_dict", "iter", "infos".

The other part that I'm not sure if is correct is this

self.action_space = Box(low=-10.0, high=10.0, shape=(12,))

# instantiate state-preprocessor and the policy
self.observation_space = Box(low=-10.0, high=10.0, shape=(60,), seed=42) #Dict( {"robot-state": Box(-1, 1, shape=(53,))} , seed=42)

self.state_preprocessor = RunningStandardScaler(self.observation_space, device=device)
self.policy = SharedModel(self.observation_space, self.action_space, device=device)

# load checkpoints
#self.state_preprocessor.load_state_dict(torch.load(statePreprocessor, weights_only=True))
self.policy.load(DNNmodel)  # same as policy.load_state_dict(torch.load(".pt"))

Then to run inference the model I do this

actions = self.policy.act(self.observation_space, role="policy")

is this correct? if yes, How can I populate the observation_space with reading from the sensors?