Equivariance test fails on C8 for EquivariantObsEnc, EquivariantResEncoder76Cyclic, and full model

[htrbao]

Description

I tested equivariance of two components from this repo (EquivariantObsEnc and EquivariantResEncoder76Cyclic) as well as the full policy model, under the cyclic group C8.
The results suggest that none of them are strictly equivariant:

• EquivariantObsEnc: fails for all rotations.
• EquivariantResEncoder76Cyclic: passes for identity (0 rotation) but fails for all other rotations.
• DiffusionEquiUNetCNNEncPolicy: fails for all rotations.

---

Test Method

I used the standard equivariance check:

• f(Tx) = apply transformation T to the input, then pass through the network
• T(f(x)) = pass input through the network, then transform the output
• Compare whether f(Tx) ≈ T(f(x)) for each group element

---

Results

EquivariantObsEnc

False 0[2pi/8] tensor(0.2026, device='cuda:1')
False 1[2pi/8] tensor(0.2031, device='cuda:1')
False 2[2pi/8] tensor(0.2006, device='cuda:1')
False 3[2pi/8] tensor(0.2044, device='cuda:1')
False 4[2pi/8] tensor(0.2044, device='cuda:1')
False 5[2pi/8] tensor(0.2025, device='cuda:1')
False 6[2pi/8] tensor(0.2030, device='cuda:1')
False 7[2pi/8] tensor(0.2027, device='cuda:1')

EquivariantResEncoder76Cyclic

True  0[2pi/8] tensor(4.1376e-18, device='cuda:1')
False 1[2pi/8] tensor(0.4769, device='cuda:1')
False 2[2pi/8] tensor(0.2295, device='cuda:1')
False 3[2pi/8] tensor(0.4780, device='cuda:1')
False 4[2pi/8] tensor(0.2746, device='cuda:1')
False 5[2pi/8] tensor(0.4778, device='cuda:1')
False 6[2pi/8] tensor(0.2294, device='cuda:1')
False 7[2pi/8] tensor(0.4769, device='cuda:1')

DiffusionEquiUNetCNNEncPolicy

False 0[2pi/8] tensor(0.6322)
False 1[2pi/8] tensor(0.6452)
False 2[2pi/8] tensor(0.6487)
False 3[2pi/8] tensor(0.6303)
False 4[2pi/8] tensor(0.6254)
False 5[2pi/8] tensor(0.6534)
False 6[2pi/8] tensor(0.6494)
False 7[2pi/8] tensor(0.6729)

---

Test Script

• The script here for checking equivariant property of EquivariantObsEnc module.

• The script here for checking equivariant property of EquivariantResEncoder76Cyclic module.

• The script here for checking equivariant property of DiffusionEquiUNetCNNEncPolicy.

---

Expected Behavior

For a properly equivariant module, all rotations in C8 should satisfy:

f(T_g x) ≈ T_g f(x)

within numerical tolerance. That means torch.allclose(y_new, y_transformed, atol=1e-4) should be True for all group elements g.

---

Actual Behavior

• EquivariantObsEnc: fails for all C8 rotations.
• EquivariantResEncoder76Cyclic: only passes for identity rotation.
• DiffusionEquiUNetCNNEncPolicy: fails for all rotations.

---

Questions

1. Are these components expected to be strictly C8-equivariant?
2. Is any additional configuration (e.g., padding, stride, field types) required to preserve equivariance?
3. If my understanding of the equivariance test is incorrect, is there another recommended way to verify the model’s equivariance property?

[pointW]

Hi @htrbao , checking the equivariance error is a great practice. However, in your code (https://github.com/htrbao/equidiff/blob/8f880f7dcee3e854df03c6add08a4a345ba932c7/EquivariantResEncoder76Cyclic.py#L23C1-L24C1), you used x = torch.randn(batch_size, 3, 512, 512).cuda(1), while this encoder is particularly designed for a 76x76 input size. You should see a lower equivariance error if you change the input size to 76x76. Moreover, I wouldn't recommend using random inputs to perform the equivariance check, as the transform function inside escnn is not perfect when it comes to images. Using something like x = torch.zeros(batch_size, 3, 76, 76).cuda(1); x[:, :, 0:10, 0:10] = 1 is probably better.

I did not get a chance to look closely at your other scripts, but I recommend you check https://github.com/pointW/equidiff/blob/main/equi_diffpo/model/vision/rot_randomizer.py to make sure you are transforming all the observation components correctly.

For your questions:

1. The models are defined under C8 group, but if your input contains images, it can only be strictly C4 equivariant because of the aliasing of image rotation.
2. padding, stride, and field type won't typically affect equivariance, but pooling might. You need to make sure you are not applying pooling in a non-divisible case, e.g., applying 2×2 pooling on a 9×9 feature map. Another thing to be careful about is the random crop augmentation. You can swap it with a center crop to make sure it won't break the symmetry.
3. Like I mentioned earlier, referring to the augmentation code when you are checking the equivariance error would be helpful.

[htrbao]

Thanks @pointW a lot for the detailed explanation! That clarifies many points.

I re-ran the equivariance test for EquivariantResEncoder76Cyclic using the correct 76×76 input and a structured input (instead of random noise), and the results now match your claim — the error is much lower and being C4 equivariant.

However, when I train and evaluate the entire model on the stack_d1 dataset, the equivariance property still seems to fail on both C4 and C8 rotations. I run it with:

python train.py --config-name=test_equi_diffusion_unet_abs task_name=stack_d1 n_demo=100

Here are the references in the repo:

• TestEquivariantWorkspace
• Config

Could you help me understand why the full model remains non-equivariant in this case?
Is it expected due to the dataset (e.g., augmentations or observation components), or is there something else I might be missing?

[pointW]

1. predict_action is a stochastic process. To check the equivariance error, you should isolate the diffusion unet and send it rotated denoising condition vector and noisy action vector. If you run predict_action as a whole, the random sampling in the denoising process will break the symmetry.
2. In https://github.com/htrbao/equidiff/blob/167d4d14ba5546ac68895411bafb7891ebfb7e48/equi_diffpo/workspace/test_equi_workspace.py#L125, [0, 1] is a irrep(1) vector (the x, y position); [3, 6], [4, 7], [5, 8] are the three other pairs of irrep(1) vectors (the three column vectors in rotation 6D).

[htrbao]

Thanks, @pointW ! After trying your suggestion, we found that EquivariantObsEnc behaves as an equivariant module. For the whole model, the equivariance error is around 0.06 for other angles, and perfectly equivariant at 0°.

Would you consider this small enough to qualify as equivariant? More generally, what error threshold do you recommend when evaluating whether a model is sufficiently equivariant?

[pointW]

0.06 is acceptable, especially if you will get significantly larger numbers when the input and output are not transformed equivariantly. I typically use 1e-3 or even smaller though.