diff --git a/packages/@aws-cdk/aws-stepfunctions-tasks/test/integ.invoke-function.expected.json b/packages/@aws-cdk/aws-stepfunctions-tasks/test/integ.invoke-function.expected.json
index 18f710ef22d77..16f8c78de01ff 100644
--- a/packages/@aws-cdk/aws-stepfunctions-tasks/test/integ.invoke-function.expected.json
+++ b/packages/@aws-cdk/aws-stepfunctions-tasks/test/integ.invoke-function.expected.json
@@ -272,7 +272,7 @@
               {
                 "Ref": "CallbackHandler4434C38D"
               },
-              "\",\"Payload\":{\"token\":\"$$.Task.Token\"}},\"Type\":\"Task\",\"Resource\":\"arn:aws:states:::lambda:invoke.waitForTaskToken\",\"ResultPath\":\"$.status\"},\"Job Complete?\":{\"Type\":\"Choice\",\"Choices\":[{\"Variable\":\"$.status\",\"StringEquals\":\"FAILED\",\"Next\":\"Job Failed\"},{\"Variable\":\"$.status\",\"StringEquals\":\"SUCCEEDED\",\"Next\":\"Final step\"}]},\"Job Failed\":{\"Type\":\"Fail\",\"Error\":\"DescribeJob returned FAILED\",\"Cause\":\"AWS Batch Job Failed\"},\"Final step\":{\"Type\":\"Pass\",\"End\":true}},\"TimeoutSeconds\":30}"
+              "\",\"Payload\":{\"token.$\":\"$$.Task.Token\"}},\"Type\":\"Task\",\"Resource\":\"arn:aws:states:::lambda:invoke.waitForTaskToken\",\"ResultPath\":\"$.status\"},\"Job Complete?\":{\"Type\":\"Choice\",\"Choices\":[{\"Variable\":\"$.status\",\"StringEquals\":\"FAILED\",\"Next\":\"Job Failed\"},{\"Variable\":\"$.status\",\"StringEquals\":\"SUCCEEDED\",\"Next\":\"Final step\"}]},\"Job Failed\":{\"Type\":\"Fail\",\"Error\":\"DescribeJob returned FAILED\",\"Cause\":\"AWS Batch Job Failed\"},\"Final step\":{\"Type\":\"Pass\",\"End\":true}},\"TimeoutSeconds\":30}"
             ]
           ]
         },
diff --git a/packages/@aws-cdk/aws-stepfunctions-tasks/test/invoke-function.test.ts b/packages/@aws-cdk/aws-stepfunctions-tasks/test/invoke-function.test.ts
index 611f30d627815..d9b3c8664af1d 100644
--- a/packages/@aws-cdk/aws-stepfunctions-tasks/test/invoke-function.test.ts
+++ b/packages/@aws-cdk/aws-stepfunctions-tasks/test/invoke-function.test.ts
@@ -15,7 +15,7 @@ beforeEach(() => {
   });
 });
 
-test('Lambda function can be used in a Task', () => {
+test('Invoke lambda with function ARN', () => {
   // WHEN
   const task = new sfn.Task(stack, 'Task', { task: new tasks.InvokeFunction(fn) });
   new sfn.StateMachine(stack, 'SM', {
@@ -39,7 +39,7 @@ test('Lambda function payload ends up in Parameters', () => {
     definition: new sfn.Task(stack, 'Task', {
       task: new tasks.InvokeFunction(fn, {
         payload: {
-          foo: 'bar'
+          foo: sfn.Data.stringAt('$.bar')
         }
       })
     })
@@ -48,45 +48,10 @@ test('Lambda function payload ends up in Parameters', () => {
   expect(stack).toHaveResource('AWS::StepFunctions::StateMachine', {
     DefinitionString: {
       "Fn::Join": ["", [
-        "{\"StartAt\":\"Task\",\"States\":{\"Task\":{\"End\":true,\"Parameters\":{\"foo\":\"bar\"},\"Type\":\"Task\",\"Resource\":\"",
+        "{\"StartAt\":\"Task\",\"States\":{\"Task\":{\"End\":true,\"Parameters\":{\"foo.$\":\"$.bar\"},\"Type\":\"Task\",\"Resource\":\"",
         { "Fn::GetAtt": ["Fn9270CBC0", "Arn"] },
         "\"}}}"
       ]]
     },
   });
 });
-
-test('Lambda function can be used in a Task with Task Token', () => {
-  const task = new sfn.Task(stack, 'Task', {
-    task: new tasks.RunLambdaTask(fn, {
-      waitForTaskToken: true,
-      payload: {
-        token: sfn.Context.taskToken
-      }
-    })
-  });
-  new sfn.StateMachine(stack, 'SM', {
-    definition: task
-  });
-
-  // THEN
-  expect(stack).toHaveResource('AWS::StepFunctions::StateMachine', {
-    DefinitionString: {
-      "Fn::Join": ["", [
-          "{\"StartAt\":\"Task\",\"States\":{\"Task\":{\"End\":true,\"Parameters\":{\"FunctionName\":\"",
-          { Ref: "Fn9270CBC0" },
-          "\",\"Payload\":{\"token\":\"$$.Task.Token\"}},\"Type\":\"Task\",\"Resource\":\"arn:aws:states:::lambda:invoke.waitForTaskToken\"}}}"
-      ]]
-    },
-  });
-});
-
-test('Task throws if waitForTaskToken is supplied but task token is not included', () => {
-  expect(() => {
-    new sfn.Task(stack, 'Task', {
-      task: new tasks.RunLambdaTask(fn, {
-        waitForTaskToken: true
-      })
-    });
-  }).toThrow(/Task Token is missing in payload/i);
-});
\ No newline at end of file
diff --git a/packages/@aws-cdk/aws-stepfunctions-tasks/test/run-lambda-task.test.ts b/packages/@aws-cdk/aws-stepfunctions-tasks/test/run-lambda-task.test.ts
new file mode 100644
index 0000000000000..545158b911768
--- /dev/null
+++ b/packages/@aws-cdk/aws-stepfunctions-tasks/test/run-lambda-task.test.ts
@@ -0,0 +1,73 @@
+import '@aws-cdk/assert/jest';
+import lambda = require('@aws-cdk/aws-lambda');
+import sfn = require('@aws-cdk/aws-stepfunctions');
+import { Stack } from '@aws-cdk/cdk';
+import tasks = require('../lib');
+
+let stack: Stack;
+let fn: lambda.Function;
+beforeEach(() => {
+  stack = new Stack();
+  fn = new lambda.Function(stack, 'Fn', {
+    code: lambda.Code.inline('hello'),
+    handler: 'index.hello',
+    runtime: lambda.Runtime.Python27,
+  });
+});
+
+test('Invoke lambda with default magic ARN', () => {
+  const task = new sfn.Task(stack, 'Task', {
+    task: new tasks.RunLambdaTask(fn, {
+      payload: {
+        foo: 'bar'
+      }
+    })
+  });
+  new sfn.StateMachine(stack, 'SM', {
+    definition: task
+  });
+
+  expect(stack).toHaveResource('AWS::StepFunctions::StateMachine', {
+    DefinitionString: {
+      "Fn::Join": ["", [
+          "{\"StartAt\":\"Task\",\"States\":{\"Task\":{\"End\":true,\"Parameters\":{\"FunctionName\":\"",
+          { Ref: "Fn9270CBC0" },
+          "\",\"Payload\":{\"foo\":\"bar\"}},\"Type\":\"Task\",\"Resource\":\"arn:aws:states:::lambda:invoke\"}}}"
+      ]]
+    },
+  });
+});
+
+test('Lambda function can be used in a Task with Task Token', () => {
+  const task = new sfn.Task(stack, 'Task', {
+    task: new tasks.RunLambdaTask(fn, {
+      waitForTaskToken: true,
+      payload: {
+        token: sfn.Context.taskToken
+      }
+    })
+  });
+  new sfn.StateMachine(stack, 'SM', {
+    definition: task
+  });
+
+  expect(stack).toHaveResource('AWS::StepFunctions::StateMachine', {
+    DefinitionString: {
+      "Fn::Join": ["", [
+          "{\"StartAt\":\"Task\",\"States\":{\"Task\":{\"End\":true,\"Parameters\":{\"FunctionName\":\"",
+          { Ref: "Fn9270CBC0" },
+          "\",\"Payload\":{\"token.$\":\"$$.Task.Token\"}},\"Type\":\"Task\",\"Resource\":\"arn:aws:states:::lambda:invoke.waitForTaskToken\"}}}"
+      ]]
+    },
+  });
+});
+
+test('Task throws if waitForTaskToken is supplied but task token is not included', () => {
+  expect(() => {
+    new sfn.Task(stack, 'Task', {
+      task: new tasks.RunLambdaTask(fn, {
+        waitForTaskToken: true
+      })
+    });
+  }).toThrow(/Task Token is missing in payload/i);
+});
\ No newline at end of file
diff --git a/packages/@aws-cdk/aws-stepfunctions/lib/json-path.ts b/packages/@aws-cdk/aws-stepfunctions/lib/json-path.ts
index ab31ef372c37c..e9f07fa2f7a4a 100644
--- a/packages/@aws-cdk/aws-stepfunctions/lib/json-path.ts
+++ b/packages/@aws-cdk/aws-stepfunctions/lib/json-path.ts
@@ -36,7 +36,8 @@ export function renderObject(obj: object | undefined): object | undefined {
   return recurseObject(obj, {
     handleString: renderString,
     handleList: renderStringList,
-    handleNumber: renderNumber
+    handleNumber: renderNumber,
+    handleBoolean: renderBoolean,
   });
 }
 
@@ -63,6 +64,10 @@ export function findReferencedPaths(obj: object | undefined): Set<string> {
       const path = jsonPathNumber(x);
       if (path !== undefined) { found.add(path); }
       return {};
+    },
+
+    handleBoolean(_key: string, _x: boolean) {
+      return {};
     }
   });
 
@@ -73,6 +78,7 @@ interface FieldHandlers {
   handleString(key: string, x: string): {[key: string]: string};
   handleList(key: string, x: string[]): {[key: string]: string[] | string };
   handleNumber(key: string, x: number): {[key: string]: number | string};
+  handleBoolean(key: string, x: boolean): {[key: string]: boolean};
 }
 
 export function recurseObject(obj: object | undefined, handlers: FieldHandlers): object | undefined {
@@ -86,6 +92,8 @@ export function recurseObject(obj: object | undefined, handlers: FieldHandlers):
       Object.assign(ret, handlers.handleNumber(key, value));
     } else if (Array.isArray(value)) {
       Object.assign(ret, recurseArray(key, value, handlers));
+    } else if (typeof value === 'boolean') {
+      Object.assign(ret, handlers.handleBoolean(key, value));
     } else if (value === null || value === undefined) {
       // Nothing
     } else if (typeof value === 'object') {
@@ -144,7 +152,7 @@ function renderString(key: string, value: string): {[key: string]: string} {
 }
 
 /**
- * Render a parameter string
+ * Render a parameter string list
  *
  * If the string value starts with '$.', render it as a path string, otherwise as a direct string.
  */
@@ -158,7 +166,7 @@ function renderStringList(key: string, value: string[]): {[key: string]: string[
 }
 
 /**
- * Render a parameter string
+ * Render a parameter number
  *
  * If the string value starts with '$.', render it as a path string, otherwise as a direct string.
  */
@@ -171,6 +179,13 @@ function renderNumber(key: string, value: number): {[key: string]: number | stri
   }
 }
 
+/**
+ * Render a parameter boolean
+ */
+function renderBoolean(key: string, value: boolean): {[key: string]: boolean} {
+    return { [key]: value };
+}
+
 /**
  * If the indicated string is an encoded JSON path, return the path
  *
diff --git a/packages/@aws-cdk/aws-stepfunctions/lib/state-graph.ts b/packages/@aws-cdk/aws-stepfunctions/lib/state-graph.ts
index 50d368493ce63..e7b382f7dbfec 100644
--- a/packages/@aws-cdk/aws-stepfunctions/lib/state-graph.ts
+++ b/packages/@aws-cdk/aws-stepfunctions/lib/state-graph.ts
@@ -1,4 +1,5 @@
 import iam = require('@aws-cdk/aws-iam');
+import { FieldUtils } from "./fields";
 import { State } from "./states/state";
 
 /**
@@ -101,7 +102,7 @@ export class StateGraph {
     public toGraphJson(): object {
         const states: any = {};
         for (const state of this.allStates) {
-            states[state.stateId] = state.toStateJson();
+            states[state.stateId] = FieldUtils.renderObject(state.toStateJson());
         }
 
         return {
diff --git a/packages/@aws-cdk/aws-stepfunctions/test/test.fields.ts b/packages/@aws-cdk/aws-stepfunctions/test/test.fields.ts
index 737ad04d43ec3..575b0d33d72e7 100644
--- a/packages/@aws-cdk/aws-stepfunctions/test/test.fields.ts
+++ b/packages/@aws-cdk/aws-stepfunctions/test/test.fields.ts
@@ -4,6 +4,8 @@ import { Context, Data, FieldUtils } from "../lib";
 export = {
   'deep replace correctly handles fields in arrays'(test: Test) {
     test.deepEqual(FieldUtils.renderObject({
+      unknown: undefined,
+      bool: true,
       literal: 'literal',
       field: Data.stringAt('$.stringField'),
       listField: Data.listAt('$.listField'),
@@ -14,6 +16,7 @@ export = {
         }
       ]
     }), {
+      'bool': true,
       'literal': 'literal',
       'field.$': '$.stringField',
       'listField.$': '$.listField',
@@ -33,10 +36,12 @@ export = {
       str: Context.stringAt('$$.Execution.StartTime'),
       count: Context.numberAt('$$.State.RetryCount'),
       token: Context.taskToken,
+      entire: Context.entireContext
     }), {
       'str.$': '$$.Execution.StartTime',
       'count.$': '$$.State.RetryCount',
-      'token.$': '$$.Task.Token'
+      'token.$': '$$.Task.Token',
+      'entire.$': '$$'
     });
 
     test.done();
@@ -44,6 +49,7 @@ export = {
 
   'find all referenced paths'(test: Test) {
     test.deepEqual(FieldUtils.findReferencedPaths({
+      bool: false,
       literal: 'literal',
       field: Data.stringAt('$.stringField'),
       listField: Data.listAt('$.listField'),