Error when attempting to transfer a policy to real world

[giangdao1402]

Hi @Toni-SM, thanks for your great work !
I try to deploy a policy from Isaac Lab to real world by imitating your code with iiwa but there is a error

self._rnn_final_states["policy"] if self.policy is self.value else outputs.get("rnn", [])
UnboundLocalError: local variable 'outputs' referenced before assignment

Here is the code

import torch
import torch.nn as nn

from skrl.agents.torch.ppo import PPO_RNN as PPO
from skrl.agents.torch.ppo import PPO_DEFAULT_CONFIG
from skrl.envs.torch import wrap_env

# Import the skrl components to build the RL system
from skrl.models.torch import GaussianMixin, Model
from skrl.resources.preprocessors.torch import RunningStandardScaler
from skrl.trainers.torch import SequentialTrainer
from skrl.utils.spaces.torch import unflatten_tensorized_space
from torch.nn import TransformerEncoder, TransformerEncoderLayer

# Define only the policy for evaluation
class Policy(GaussianMixin, Model):
    def __init__(self, observation_space, action_space, device, clip_actions=False,
                 clip_log_std=True, min_log_std=-20, max_log_std=2, reduction="sum",
                 num_envs=1, num_layers=1, hidden_size=64, sequence_length=128):
        Model.__init__(self, observation_space, action_space, device)
        GaussianMixin.__init__(self, clip_actions, clip_log_std, min_log_std, max_log_std, reduction)

        self.num_envs = num_envs
        self.num_layers = num_layers
        self.hidden_size = hidden_size
        self.sequence_length = sequence_length
        self.embed_dim = 128

        # =============== 1. CNN ===============
        self.cnn_backbone = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=8, stride=4),
            nn.ReLU(),
            nn.Conv2d(32, 64, kernel_size=4, stride=2),
            nn.ReLU(),
            nn.Conv2d(64, 64, kernel_size=3, stride=1),
            nn.ReLU(),
        )
        self.num_patches = 7 * 7 

        # =============== 2. ViT ===============
        self.patch_embed = nn.Linear(64, self.embed_dim)
        self.pos_embed = nn.Parameter(torch.zeros(1, self.num_patches, self.embed_dim))
        self.vit_encoder = TransformerEncoder(
            TransformerEncoderLayer(d_model=self.embed_dim, nhead=4, dim_feedforward=256, batch_first=True),
            num_layers=2
        )

        # =============== 3. State encoder ===============
        # state_dim = observation_space["robot-state"].shape[0]
        state_dim = observation_space["policy"]["robot-state"].shape[0]
        self.state_encoder = nn.Sequential(
            nn.Linear(state_dim, 128),
            nn.ELU(),
            nn.Linear(128, self.embed_dim),
            nn.ELU()
        )

        # =============== 4. Cross-Attention ===============
        self.cross_attn = nn.MultiheadAttention(embed_dim=self.embed_dim, num_heads=4, batch_first=True)
        self.norm1 = nn.LayerNorm(self.embed_dim)
        self.norm2 = nn.LayerNorm(self.embed_dim)
        self.fusion_mlp = nn.Sequential(
            nn.Linear(self.embed_dim, 256),
            nn.ELU(),
            nn.Linear(256, 128),   # ← output 128D
            nn.ELU()
        )

        # =============== 5. LSTM ===============
        self.lstm = nn.LSTM(
            input_size=128,          
            hidden_size=hidden_size,
            num_layers=num_layers,
            batch_first=True
        )

        # =============== 6. POLICY HEAD ===============
        self.net = nn.Sequential(
            nn.Linear(hidden_size, 256),
            nn.ELU(),
            nn.Linear(256, 128),
            nn.ELU(),
            nn.Linear(128, self.num_actions)  
        )

        self.log_std_parameter = nn.Parameter(torch.zeros(self.num_actions))

    def get_specification(self):
        return {
            "rnn": {
                "sequence_length": self.sequence_length,
                "sizes": [
                    (self.num_layers, self.num_envs, self.hidden_size),
                    (self.num_layers, self.num_envs, self.hidden_size)
                ]
            }
        }

    def compute(self, inputs, role):
        states = unflatten_tensorized_space(self.observation_space, inputs["states"])
        if isinstance(states, dict) and "policy" in states:
            states = states["policy"]
        terminated = inputs.get("terminated", None)
        hidden_states, cell_states = inputs["rnn"][0], inputs["rnn"][1]

        camera = states["camera"]
        if camera.dim() == 3:
            camera = camera.unsqueeze(-1)
        assert camera.shape[-1] == 1, "Depth must be 1 channel"
        camera = camera.permute(0, 3, 1, 2)

        # CNN + ViT
        cnn_feat = self.cnn_backbone(camera)
        patches = cnn_feat.flatten(2).permute(0, 2, 1)
        patches = self.patch_embed(patches) + self.pos_embed
        image_features = self.vit_encoder(patches)

        # State encoder
        state_emb = self.state_encoder(states["robot-state"])
        state_query = state_emb.unsqueeze(1)

        # Cross-attention
        attn_out, _ = self.cross_attn(query=state_query, key=image_features, value=image_features)
        attn_out = attn_out.squeeze(1)
        x = self.norm1(state_emb + attn_out)
        fused = self.norm2(x + self.fusion_mlp(x))  # ← (B, 128)

        if self.training:
            rnn_input = fused.view(-1, self.sequence_length, fused.shape[-1])
            hidden_states = hidden_states.view(self.num_layers, -1, self.sequence_length, hidden_states.shape[-1])
            cell_states = cell_states.view(self.num_layers, -1, self.sequence_length, cell_states.shape[-1])
            hidden_states = hidden_states[:, :, 0, :].contiguous()
            cell_states = cell_states[:, :, 0, :].contiguous()

            if terminated is not None and torch.any(terminated):
                rnn_outputs = []
                terminated = terminated.view(-1, self.sequence_length)
                indexes = [0] + (terminated[:, :-1].any(dim=0).nonzero(as_tuple=True)[0] + 1).tolist() + [self.sequence_length]
                for i in range(len(indexes) - 1):
                    i0, i1 = indexes[i], indexes[i + 1]
                    rnn_output, (hidden_states, cell_states) = self.lstm(rnn_input[:, i0:i1, :], (hidden_states, cell_states))
                    hidden_states[:, (terminated[:, i1 - 1]), :] = 0
                    cell_states[:, (terminated[:, i1 - 1]), :] = 0
                    rnn_outputs.append(rnn_output)
                rnn_output = torch.cat(rnn_outputs, dim=1)
                rnn_states = (hidden_states, cell_states)
            else:
                rnn_output, rnn_states = self.lstm(rnn_input, (hidden_states, cell_states))
        else:
            rnn_input = fused.view(-1, 1, fused.shape[-1])
            rnn_output, rnn_states = self.lstm(rnn_input, (hidden_states, cell_states))

        rnn_output = torch.flatten(rnn_output, start_dim=0, end_dim=1)  # (N*L, hidden_size)
        print("rnn_output: ", rnn_output)

        actions = torch.tanh(self.net(rnn_output))  
        return actions, self.log_std_parameter, {"rnn": [rnn_states[0], rnn_states[1]]}


from nav_ros2_env import MobileNavEnvROS2

env = MobileNavEnvROS2()

# wrap the environment
env = wrap_env(env)

device = env.device


# Instantiate the agent's policy.
# PPO requires 2 models, visit its documentation for more details
# https://skrl.readthedocs.io/en/latest/modules/skrl.agents.ppo.html#spaces-and-models
models_ppo = {}
models_ppo["policy"] = Policy(env.observation_space, env.action_space, device, num_envs=1)


# Configure and instantiate the agent.
# Only modify some of the default configuration, visit its documentation to see all the options
# https://skrl.readthedocs.io/en/latest/modules/skrl.agents.ppo.html#configuration-and-hyperparameters
cfg_ppo = PPO_DEFAULT_CONFIG.copy()
cfg_ppo["random_timesteps"] = 0
cfg_ppo["learning_starts"] = 0
cfg_ppo["state_preprocessor"] = RunningStandardScaler
cfg_ppo["state_preprocessor_kwargs"] = {"size": env.observation_space, "device": device}
# logging to TensorBoard each 32 timesteps an ignore checkpoints
cfg_ppo["experiment"]["write_interval"] = 32
cfg_ppo["experiment"]["checkpoint_interval"] = 0

agent = PPO(
    models=models_ppo,
    memory=None,
    cfg=cfg_ppo,
    observation_space=env.observation_space,
    action_space=env.action_space,
    device=device,
)

# load checkpoints
agent.load("/home/cei/giang_ws/RLNav-IsaacLab-main/agent_50000.pt")



# Configure and instantiate the RL trainer
cfg_trainer = {"timesteps": 1000, "headless": True}
trainer = SequentialTrainer(cfg=cfg_trainer, env=env, agents=agent)

# start evaluation
trainer.eval()

Could you suggest me the solution ? Thank you !