TypeError: 'int' object is not subscriptable in locomotion branch

[ziwon-park]

Bug: TypeError in locomotion branch safety constraints

Error: TypeError: 'int' object is not subscriptable in diffusion.py:623
A TypeError: 'int' object is not subscriptable occurs when running the locomotion branch during the sampling phase.

Environment

• Branch : locomotion
• File : diffuser/models/diffusion.py
• Python version : 3.8

Full Traceback

Traceback (most recent call last):
  File "scripts/train.py", line 124, in <module>
    trainer.train(n_train_steps=args.n_steps_per_epoch)
  File "/home/work/CBF-CFM/ziwon/SafeDiffuser/diffuser/utils/training.py", line 132, in train
    self.render_samples()
  File "/home/work/CBF-CFM/ziwon/SafeDiffuser/diffuser/utils/training.py", line 208, in render_samples
    samples = self.ema_model(conditions)
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1423, in *call*impl
    return forward_call(*input, **kwargs)
  File "/home/work/CBF-CFM/ziwon/SafeDiffuser/diffuser/models/diffusion.py", line 851, in forward
    return self.conditional_sample(cond, *args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/home/work/CBF-CFM/ziwon/SafeDiffuser/diffuser/models/diffusion.py", line 230, in conditional_sample
    return self.p_sample_loop(shape, cond, return_chain = True, **sample_kwargs)    # debug
  File "/usr/local/lib/python3.8/dist-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/home/work/CBF-CFM/ziwon/SafeDiffuser/diffuser/models/diffusion.py", line 194, in p_sample_loop
    x, b_min = self.invariance_cheetah(x_t, x)
  File "/usr/local/lib/python3.8/dist-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/home/work/CBF-CFM/ziwon/SafeDiffuser/diffuser/models/diffusion.py", line 623, in invariance_cheetah
    radius = (radius - self.mean[0]) / self.std[0]
TypeError: 'int' object is not subscriptable

Problem

In the GaussianDiffusion class constructor (diffusion.py line ~75), self.mean and self.std are initialized as integers:

# Constructor initializes as int
self.mean = 0
self.std = 0

However, the safety constraint methods (invariance_cheetah, invariance_hopper, etc.) expect these to be arrays/tensors for normalization:

def invariance_cheetah(self, x, xp1):
    # ...
    radius = (radius - self.mean[0]) / self.std[0]  # ← Tries to index integer
    cx = (cx - self.mean[14]) / self.std[14]        # ← Same issue
    cy = (cy - self.mean[0]) / self.std[0]          # ← Same issue

Reproduce

1. Switch to locomotion branch
2. Run training with cheetah environment (python scripts/train.py --dataset halfcheetah-medium-expert-v2)
3. Error occurs during sampling

Affected Methods

All safety constraint methods: invariance*, invariance_*_cf, invariance_*_cpx

Naively editing just self.mean and self.std like

        self.transition_dim = observation_dim + action_dim
        self.mean = torch.zeros(self.transition_dim)
        self.std = torch.ones(self.transition_dim)

causes another errors 😂 Would appreciate any insights or suggestions on the best approach to fix this.

Thank you for taking the time to look into this issue.

[Weixy21]

You may just set self.mean = torch.tensor(0) self.std = torch.tensor(0) to fix this problem or simply remove/comment out these two lines during training as they are only used during inference for now.

I will update the code later.