Depends on the circumstances. If someone gets a life sentence, then it's probably not a benefit to society to work on rehabilitation. If someone gets 5 years or so for a non-violent offense, then perhaps providing resources that foster hope of permanent change could benefit society rather than derail them from society entirely just to punish them. That can be a balance many factors, such as a display of discipline through good behavior in jail/prison, work release, learning a skill/trade so one can get a job to support themselves (the primary precursor to being a benefit to society), and some sort of job placement assistance or at least the removal of so many barriers.

Some of these exist to some degree already. Some of the issues are systemic and not just related to those in prison. For example, the decline of real wages over the past 50 years and the rise of barriers to entry for many things (renting, jobs, etc) has made it difficult for lower skilled individuals to succeed. If you feel you have no hope in the system, then you will be more inclined to act outside the system.

None of this is an excuse for someone to break the law, merely an observation that the current system might not be maximizing societal benefits but rather too focused on punishment alone.