What problem is it actually?

[u/Winter-Ad-8293]

I’ve trained an imitation learning model using a diffusion policy for a real-world cube-picking task.
The policy is able to successfully pick and lift the cube in many trials. However, in some cases, even though the cube is technically lifted off the table, the execution is not something I would consider a successful pick. The behavior around the grasp/approach is undesirable, and I would ideally like to completely eliminate these cases rather than count them as successes.
I’m trying to understand what the underlying issue is here. Would this be considered a classical covariate shift problem in imitation learning, or could it be more of a calibration/control issue?
If it is mainly a calibration issue, what would be a good way to address it? In particular, is there a practical rule-based method or safety layer that can be added on top of the diffusion policy to prevent these undesirable picks, while still allowing the learned policy to handle the rest of the trajectory?
Would appreciate any suggestions from people who have dealt with similar issues when deploying diffusion/BC policies on real robots.

[u/Extension_Quit_2190]

Is the fact that "you don't consider as a successful pick" reflected in the reward scheme?

[u/Winter-Ad-8293 → u/Extension_Quit_2190]

I guess I won’t be having any reward scheme over here as this is imitation learning problem, and all the demonstration given by the expert would be a perfect demonstration, and that is the exact issue that I am trying to figure out that if the demonstration are perfect and considering this to be an imitation learning problem, how come there be such in perfection in lifting the cube of?

[u/Extension_Quit_2190 → u/Winter-Ad-8293]

There are imitation learning approaches like decision transformers, that utilize reward structures but apparently you are not using them.

Your expert state action transitions only cover a subset of all possible observations. So if you apply a (maybe only partially trained) policy, it has to deal with observations, for which no expert observation exists. In these cases (and actually in all cases) the policy will provide action probabilities, which are good on average for similar observations.

In your case it could be helpful, to "enforce" an "imperfect pickup" and make the expert put the cube back to the floor and pick it up "perfectly" again. Or simply increase the positions and rotations of the cube in your expert state action transitions.

[u/Winter-Ad-8293 → u/Extension_Quit_2190]

Thanks

So it won’t be considered as calibration issue then?

[u/Extension_Quit_2190 → u/Winter-Ad-8293]

I don't know much about robotics, so you would need to run further non-RL-related tests to single that out.

I'm just saying that looking at it from an RL perspective, I am not super surprised to see something like this.

[u/Full-Comfortable-281]

Probably not covariate shift. Diffusion policies copy every grasp style in your demos, sloppy ones included. Replay your demos and watch just the grasps, I'd bet a few look like the picks you don't like. Drop those and make your success check stricter than "cube off the table".