I know this reaction is not rational, but still, my first reaction was: In such environment (where it is possible to tell the difference between irrational and rational crime, and punish accordingly), becoming rational means losing your "get of out the jail once" card, and that's not fair! The more rational you are, the wider range of your possible crimes becomes punishable. You are being punished for being rational.

Maybe rational defector was the wrong way to put it. I don't mean punish people who test high on rationality, I mean punish in the cases where it's a calculated defection for personal gain. Punish in cases where tit for tat is actually an effective strategy.

Some crimes just aren't done for personal gain, and those should have alternate strategy. Of course, what the alternate strategy is is still open, and distinguishing between them is difficult, as you say:

I guess this intuition is based on the real-world situations, where the psychologists are not perfect, the justice is not perfect, and therefore any rule like this has big chance to be heavily abused. (As in: If you have a good lawyer, your crimes will be declared irrational, and you will be sentenced to two weeks of group therapy. Meanwhile the average Joe does the same thing and gets hanged.)

At our level, I don't think we are able to distinguish between crimes that should get punishment, and things where punishment is ineffective. It's just useful to understand that justice is about game theory, not revenge.