diff --git a/src/main/resources/templates/gluon/reinforcement/environment/Environment.ftl b/src/main/resources/templates/gluon/reinforcement/environment/Environment.ftl
index 0bb1adb93d0e7911f8ef96c180c9e9976c039d83..cc02e137810327dad6d8d4dcadc88a5bc169c639 100644
--- a/src/main/resources/templates/gluon/reinforcement/environment/Environment.ftl
+++ b/src/main/resources/templates/gluon/reinforcement/environment/Environment.ftl
@@ -212,7 +212,7 @@ class RosEnvironment(Environment):
         rospy.signal_shutdown('Program ended!')
 
     def __state_callback(self, data):
-        self.__last_received_state = np.array(data.data, dtype='float32')
+        self.__last_received_state = np.array(data.data, dtype='float32').reshape((<#list config.stateDim as d>${d},</#list>))
         rospy.logdebug('Received state: {}'.format(self.__last_received_state))
         self.__waiting_for_state_update = False