feat(g05): make policy runtime native

2026-08-08 17:39:44 +00:00 · 2026-07-29 16:58:02 +02:00
parent 916fc866f5
commit e6d7eea88a
9 changed files with 2372 additions and 68 deletions
@@ -0,0 +1,83 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+
+from lerobot.policies.g05.action_codec_g05 import G05NativeActionCodec, _BinarySequenceCodec
+
+
+def _tiny_codec_config() -> dict:
+    return {
+        "parts_meta": {
+            "left_control": 3,
+            "left_gripper": 1,
+            "right_control": 3,
+            "right_gripper": 1,
+        },
+        "rule_based_key_patterns": ["gripper"],
+        "rule_based_min_block_len": 1,
+        "rule_based_binarize_threshold": 0.0,
+        "num_residuals": 2,
+        "model_arch": {
+            "horizon": 8,
+            "horizon_patch_size": 2,
+            "max_component_dim": 3,
+            "conv_in_action_kernel": 2,
+            "encoder_channels": 64,
+            "c_mults": [1],
+            "strides": [[1, 1]],
+            "transformer_depths": [1],
+            "latent_dim": 16,
+            "num_heads": 1,
+            "dim_heads": 64,
+            "rope_base": 10_000,
+            "ffn_mult": 2,
+            "n_codebooks": 2,
+            "codebook_size": 16,
+            "codebook_dim": 4,
+            "use_block_dct": False,
+        },
+    }
+
+
+def test_binary_sequence_codec_roundtrip_repairs_short_middle_runs() -> None:
+    codec = _BinarySequenceCodec(sequence_length=8, min_block_length=1, vocab_size=16)
+    values = torch.tensor([[0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 1.0, 1.0]])
+
+    tokens = codec.encode(values, threshold=0.5)
+    decoded = codec.decode(tokens)
+
+    assert tokens.shape == (1, codec.num_tokens)
+    assert decoded.tolist() == [[0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0]]
+
+
+def test_native_action_codec_language_roundtrip_and_absent_groups() -> None:
+    codec = G05NativeActionCodec(_tiny_codec_config(), action_token_begin=100)
+    actions = torch.linspace(-1, 1, 8 * 8).reshape(8, 8)
+
+    token_ids = codec.encode_for_language({"value": actions})
+    decoded, absent = codec.decode_language_tokens(
+        torch.tensor(token_ids),
+        horizon=8,
+        action_dim=8,
+    )
+
+    assert len(token_ids) == codec.action_token_length
+    assert decoded.shape == (8, 8)
+    assert torch.isfinite(decoded).all()
+    assert absent == set()
+    assert all(key.startswith("model.") for key in codec.module.state_dict())
+
+    empty, absent = codec.decode_language_tokens(torch.empty(0, dtype=torch.long), horizon=8, action_dim=8)
+    assert absent == set(codec.parts)
+    assert torch.equal(empty, torch.zeros_like(empty))