diff --git a/libs/robot.lua b/libs/robot.lua
index 7f8b0f5..05794f5 100644
--- a/libs/robot.lua
+++ b/libs/robot.lua
@@ -15,11 +15,20 @@ local function identity(x)
 end
 
 local function getConfig(givenConfig)
+  if givenConfig.onSave and not givenConfig.onLoad then
+    error('robot.getConfig: custom "onSave" passed in config but no "onLoad"')
+  end
+
+  if givenConfig.onLoad and not givenConfig.onSave then
+    error('robot.getConfig: custom "onLoad" passed in config but no "onSave"')
+  end
+
   return {
     autosave = givenConfig.autosave or DEFAULT_CONFIG.autosave,
     autoload = givenConfig.autoload or DEFAULT_CONFIG.autoload,
     stateFilePath = givenConfig.stateFilePath or DEFAULT_CONFIG.stateFilePath,
-    onSave = givenConfig.onSave or identity
+    onSave = givenConfig.onSave or identity,
+    onLoad = givenConfig.onLoad or identity
   }
 end
 
@@ -28,8 +37,7 @@ local function createDefaultState()
     y = 0,
     x = 0,
     z = 0,
-    -- | BACKWARD | LEFT | RIGHT
-    dir = 'FORWARD'
+    dir = 'FORWARD' -- FORWARD | BACKWARD | LEFT | RIGHT
   }
 end
 
@@ -41,13 +49,15 @@ local saveRobotState = function(stateFile, onSave, state)
     error('saveRobotState: cannot open ' .. stateFile ..  ' file!')
   end
 
-  file.write(textutils.serializeJSON(onSave(state)))
+  local transformedState = onSave(state)
+
+  file.write(textutils.serializeJSON(transformedState))
   file.close()
 
-  return state
+  return transformedState
 end
 
-local loadRobotState = function(stateFile, onSave)
+local loadRobotState = function(stateFile, onLoad, onSave)
   local file = fs.open(stateFile, 'r')
 
   if not file then
@@ -63,7 +73,7 @@ local loadRobotState = function(stateFile, onSave)
     return robot.saveRobotState(stateFile, onSave)
   end
 
-  return parsedResult
+  return onLoad(parsedResult)
 end
 
 local getAutoSave = function(config, state)
@@ -83,7 +93,7 @@ api.create = function(state, givenConfig)
   local config = getConfig(givenConfig)
 
   if config.autoload then
-    state = loadRobotState(config.stateFilePath)
+    state = loadRobotState(config.stateFilePath, config.onLoad, config.onSave)
   end
 
   local autoSave = getAutoSave(config, state)
@@ -198,8 +208,8 @@ api.create = function(state, givenConfig)
     return saveRobotState(config.stateFilePath, config.onSave, state)
   end
 
-  robot.loadState = function(stateFile)
-    return loadRobotState(config.stateFilePath, config.onSave)
+  robot.loadState = function()
+    return loadRobotState(config.stateFilePath, config.onLoad, config.onSave)
   end
 
   return robot