# Design loss function for model-based reinforcement learning

I'm doing some model-based reinforcement learning, and I'm stuck at how to better design the loss function for fitting the dynamic model of the environment.

In continuous state and action space, the common practice is to perform a regression to predict $$S(t+1)$$ given $$S(t)$$ and $$A(t)$$ where $$S$$ and $$A$$ are states and actions respectively, and the regression uses L2 loss. But I can easily imagine that some parts of the states contribute only slightly to the overall rewards/control but are hard to predict, and one extreme example is a rat is moving randomly on my robot, and the position of the rat contributes only slightly to the overall dynamics but is hard to predict, and using L2 loss could make the model focus unproportionally on predicting the rat position. Obviously the easiest way is to drop such states but these states still have a non-zero contribution.

Is there any literature or research on how to properly design model loss function, so that the model can be better for the overall sequential decision making process?