[WIP] perf(policies): Make ACT policy compatible with torch.compile

[tc-huang]

Note

Work in Progress !
This PR is my attempt to address issue #2061. As it's a complex issue and the PR is still under development, all collaboration, help, and suggestions would be greatly appreciated!

What this does

(WIP) Explain what this PR does.

• Related to issue Make policies compatible with torch.compile #2061
• Changes to make the policy compile-compatible
• Brief description of the fix

W1018 23:02:25.828000 103342 torch/_dynamo/variables/tensor.py:913] [0/0] Graph break from `Tensor.item()`, consider setting:
W1018 23:02:25.828000 103342 torch/_dynamo/variables/tensor.py:913] [0/0]     torch._dynamo.config.capture_scalar_outputs = True
W1018 23:02:25.828000 103342 torch/_dynamo/variables/tensor.py:913] [0/0] or:
W1018 23:02:25.828000 103342 torch/_dynamo/variables/tensor.py:913] [0/0]     env TORCHDYNAMO_CAPTURE_SCALAR_OUTPUTS=1
W1018 23:02:25.828000 103342 torch/_dynamo/variables/tensor.py:913] [0/0] to include these operations in the captured graph.
W1018 23:02:25.828000 103342 torch/_dynamo/variables/tensor.py:913] [0/0] 
W1018 23:02:25.828000 103342 torch/_dynamo/variables/tensor.py:913] [0/0] Graph break: from user code at:
W1018 23:02:25.828000 103342 torch/_dynamo/variables/tensor.py:913] [0/0]   File "/home/tc-huang/Desktop/lerobot/src/lerobot/policies/act/modeling_act.py", line 147, in forward
W1018 23:02:25.828000 103342 torch/_dynamo/variables/tensor.py:913] [0/0]     loss_dict = {"l1_loss": l1_loss.item()}
W1018 23:02:25.828000 103342 torch/_dynamo/variables/tensor.py:913] [0/0] 
W1018 23:02:25.828000 103342 torch/_dynamo/variables/tensor.py:913] [0/0] 

• src/lerobot/policies/act/modeling_act.py

  -    loss_dict = {"l1_loss": l1_loss.item()}
  +    loss_dict = {"l1_loss": l1_loss}
       if self.config.use_vae:
            # Calculate Dₖₗ(latent_pdf || standard_normal). Note: After computing the KL-divergence for
            # each dimension independently, we sum over the latent dimension to get the total
            # KL-divergence per batch element, then take the mean over the batch.
            # (See App. B of https://huggingface.co/papers/1312.6114 for more details).
            mean_kld = (
                (-0.5 * (1 + log_sigma_x2_hat - mu_hat.pow(2) - (log_sigma_x2_hat).exp())).sum(-1).mean()
            )
  -         loss_dict["kld_loss"] = mean_kld.item()
  +         loss_dict["kld_loss"] = mean_kld
            loss = l1_loss + mean_kld * self.config.kl_weight
        else:
            loss = l1_loss
  
        return loss, loss_dict

• src/lerobot/scripts/lerobot_train.py

   with accelerator.autocast():
       loss, output_dict = policy.forward(batch)
  +    output_dict = {k: v.item() if isinstance(v, torch.Tensor) else v for k, v in output_dict.items()}

Current progress

• ✅ Compilation: PASSED
• ❌ Correctness: FAILED
• ❌ Performance: FAILED
• ✅ Benchmarking: PASSED

How it was tested

(WIP) Explain/show how you tested your changes.

• Baseline and final benchmark reports

- lerobot version: 0.3.4
- Platform: Linux-6.8.0-84-generic-x86_64-with-glibc2.35
- Python version: 3.10.12
- Huggingface Hub version: 0.35.1
- Datasets version: 4.1.1
- Numpy version: 2.2.6
- PyTorch version: 2.7.1+cu126
- Is PyTorch built with CUDA support?: True
- Cuda version: 12.6
- GPU model: NVIDIA GeForce RTX 4090
- Using GPU in script?: Yes

python benchmarks/policies_compilation/benchmark_inference_compile_lerobot.py --policy act --device cuda --output benchmarks/policies_compilation/baseline_act_report.md

How to checkout & try? (for the reviewer)

(WIP) Provide a simple way for the reviewer to try out your changes.

python benchmarks/policies_compilation/benchmark_inference_compile_lerobot.py --policy act --device cuda --output benchmarks/policies_compilation/baseline_act_report.md