diff --git a/flink-ml-iteration/src/main/java/org/apache/flink/iteration/proxy/state/ProxyOperatorStateBackend.java b/flink-ml-iteration/src/main/java/org/apache/flink/iteration/proxy/state/ProxyOperatorStateBackend.java
index a655886..58fc57a 100644
--- a/flink-ml-iteration/src/main/java/org/apache/flink/iteration/proxy/state/ProxyOperatorStateBackend.java
+++ b/flink-ml-iteration/src/main/java/org/apache/flink/iteration/proxy/state/ProxyOperatorStateBackend.java
@@ -21,6 +21,11 @@
 import org.apache.flink.api.common.state.ListState;
 import org.apache.flink.api.common.state.ListStateDescriptor;
 import org.apache.flink.api.common.state.MapStateDescriptor;
+import org.apache.flink.api.common.state.StateDescriptor;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.typeutils.ListTypeInfo;
+import org.apache.flink.api.java.typeutils.MapTypeInfo;
+import org.apache.flink.iteration.utils.ReflectionUtils;
 import org.apache.flink.runtime.checkpoint.CheckpointOptions;
 import org.apache.flink.runtime.state.CheckpointStreamFactory;
 import org.apache.flink.runtime.state.OperatorStateBackend;
@@ -50,20 +55,35 @@
     @Override
     public <K, V> BroadcastState<K, V> getBroadcastState(MapStateDescriptor<K, V> stateDescriptor)
             throws Exception {
-        MapStateDescriptor<K, V> newDescriptor =
-                new MapStateDescriptor<>(
-                        stateNamePrefix.prefix(stateDescriptor.getName()),
-                        stateDescriptor.getKeySerializer(),
-                        stateDescriptor.getValueSerializer());
+        MapStateDescriptor<K, V> newDescriptor;
+        if (stateDescriptor.isSerializerInitialized()) {
+            newDescriptor =
+                    new MapStateDescriptor<>(
+                            stateNamePrefix.prefix(stateDescriptor.getName()),
+                            stateDescriptor.getKeySerializer(),
+                            stateDescriptor.getValueSerializer());
+        } else {
+            MapTypeInfo<K, V> mapTypeInfo = getMapTypeInfo(stateDescriptor);
+            newDescriptor =
+                    new MapStateDescriptor<>(
+                            stateNamePrefix.prefix(stateDescriptor.getName()),
+                            mapTypeInfo.getKeyTypeInfo(),
+                            mapTypeInfo.getValueTypeInfo());
+        }
         return wrappedBackend.getBroadcastState(newDescriptor);
     }
 
     @Override
     public <S> ListState<S> getListState(ListStateDescriptor<S> stateDescriptor) throws Exception {
         ListStateDescriptor<S> newDescriptor =
-                new ListStateDescriptor<>(
-                        stateNamePrefix.prefix(stateDescriptor.getName()),
-                        stateDescriptor.getElementSerializer());
+                stateDescriptor.isSerializerInitialized()
+                        ? new ListStateDescriptor<>(
+                                stateNamePrefix.prefix(stateDescriptor.getName()),
+                                stateDescriptor.getElementSerializer())
+                        : new ListStateDescriptor<>(
+                                stateNamePrefix.prefix(stateDescriptor.getName()),
+                                getElementTypeInfo(stateDescriptor));
+
         return wrappedBackend.getListState(newDescriptor);
     }
 
@@ -71,9 +91,13 @@
     public <S> ListState<S> getUnionListState(ListStateDescriptor<S> stateDescriptor)
             throws Exception {
         ListStateDescriptor<S> newDescriptor =
-                new ListStateDescriptor<S>(
-                        stateNamePrefix.prefix(stateDescriptor.getName()),
-                        stateDescriptor.getElementSerializer());
+                stateDescriptor.isSerializerInitialized()
+                        ? new ListStateDescriptor<>(
+                                stateNamePrefix.prefix(stateDescriptor.getName()),
+                                stateDescriptor.getElementSerializer())
+                        : new ListStateDescriptor<>(
+                                stateNamePrefix.prefix(stateDescriptor.getName()),
+                                getElementTypeInfo(stateDescriptor));
         return wrappedBackend.getUnionListState(newDescriptor);
     }
 
@@ -125,4 +149,16 @@
             throws Exception {
         return wrappedBackend.snapshot(checkpointId, timestamp, streamFactory, checkpointOptions);
     }
+
+    @SuppressWarnings("unchecked,rawtypes")
+    private <S> TypeInformation<S> getElementTypeInfo(ListStateDescriptor<S> stateDescriptor) {
+        return ((ListTypeInfo)
+                        ReflectionUtils.getFieldValue(
+                                stateDescriptor, StateDescriptor.class, "typeInfo"))
+                .getElementTypeInfo();
+    }
+
+    private <K, V> MapTypeInfo<K, V> getMapTypeInfo(MapStateDescriptor<K, V> stateDescriptor) {
+        return ReflectionUtils.getFieldValue(stateDescriptor, StateDescriptor.class, "typeInfo");
+    }
 }
diff --git a/flink-ml-tests/src/test/java/org/apache/flink/test/iteration/BoundedPerRoundStreamIterationITCase.java b/flink-ml-tests/src/test/java/org/apache/flink/test/iteration/BoundedPerRoundStreamIterationITCase.java
index 825aa31..a1b5609 100644
--- a/flink-ml-tests/src/test/java/org/apache/flink/test/iteration/BoundedPerRoundStreamIterationITCase.java
+++ b/flink-ml-tests/src/test/java/org/apache/flink/test/iteration/BoundedPerRoundStreamIterationITCase.java
@@ -19,6 +19,10 @@
 package org.apache.flink.test.iteration;
 
 import org.apache.flink.api.common.functions.JoinFunction;
+import org.apache.flink.api.common.state.BroadcastState;
+import org.apache.flink.api.common.state.ListState;
+import org.apache.flink.api.common.state.ListStateDescriptor;
+import org.apache.flink.api.common.state.MapStateDescriptor;
 import org.apache.flink.api.common.typeinfo.BasicTypeInfo;
 import org.apache.flink.api.common.typeinfo.Types;
 import org.apache.flink.api.java.tuple.Tuple2;
@@ -26,12 +30,14 @@
 import org.apache.flink.iteration.IterationBody;
 import org.apache.flink.iteration.IterationBodyResult;
 import org.apache.flink.iteration.IterationConfig;
+import org.apache.flink.iteration.IterationConfig.OperatorLifeCycle;
 import org.apache.flink.iteration.Iterations;
 import org.apache.flink.iteration.ReplayableDataStreamList;
 import org.apache.flink.ml.common.datastream.EndOfStreamWindows;
 import org.apache.flink.ml.common.iteration.TerminateOnMaxIter;
 import org.apache.flink.runtime.jobgraph.JobGraph;
 import org.apache.flink.runtime.minicluster.MiniCluster;
+import org.apache.flink.runtime.state.StateInitializationContext;
 import org.apache.flink.streaming.api.datastream.DataStream;
 import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
 import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
@@ -55,6 +61,8 @@
 import org.junit.Rule;
 import org.junit.Test;
 
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Map;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.LinkedBlockingQueue;
@@ -73,6 +81,7 @@
 
     private SharedReference<BlockingQueue<OutputRecord<Integer>>> collectedOutputRecord;
     private SharedReference<BlockingQueue<Long>> collectedWatermarks;
+    private SharedReference<BlockingQueue<Long>> collectedOutputs;
 
     @Before
     public void setup() throws Exception {
@@ -81,6 +90,7 @@
 
         collectedOutputRecord = sharedObjects.add(new LinkedBlockingQueue<>());
         collectedWatermarks = sharedObjects.add(new LinkedBlockingQueue<>());
+        collectedOutputs = sharedObjects.add(new LinkedBlockingQueue<>());
     }
 
     @After
@@ -136,6 +146,33 @@
                 .forEachRemaining(x -> assertEquals(Long.MAX_VALUE, (long) x));
     }
 
+    @Test
+    public void testPerRoundIterationWithState() throws Exception {
+        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
+        env.setParallelism(1);
+        DataStream<Long> broadcastStream = env.fromElements(1L);
+        DataStream<Long> inputStream = env.fromElements(1L);
+        DataStreamList outputStream =
+                Iterations.iterateBoundedStreamsUntilTermination(
+                        DataStreamList.of(inputStream),
+                        ReplayableDataStreamList.replay(broadcastStream),
+                        IterationConfig.newBuilder()
+                                .setOperatorLifeCycle(OperatorLifeCycle.PER_ROUND)
+                                .build(),
+                        new PerRoundIterationBodyWithState());
+
+        outputStream.<Long>get(0).addSink(new LongSink(collectedOutputs));
+        JobGraph jobGraph = env.getStreamGraph().getJobGraph();
+        miniCluster.executeJobBlocking(jobGraph);
+
+        List<Long> result = new ArrayList<>(3);
+        collectedOutputs.get().drainTo(result);
+        assertEquals(3, result.size());
+        for (long value : result) {
+            assertEquals(1L, value);
+        }
+    }
+
     private static JobGraph createPerRoundJobGraph(
             int numSources,
             int numRecordsPerSource,
@@ -229,6 +266,56 @@
         }
     }
 
+    private static class PerRoundIterationBodyWithState implements IterationBody {
+
+        @Override
+        public IterationBodyResult process(
+                DataStreamList variableStreams, DataStreamList dataStreams) {
+            DataStream<Long> variableStream = variableStreams.get(0);
+
+            DataStream<Long> feedback =
+                    variableStream.transform("mapWithState", Types.LONG, new MapWithState());
+
+            DataStream<Integer> terminationCriteria =
+                    feedback.<Long>flatMap(new TerminateOnMaxIter(2)).returns(Types.INT);
+
+            return new IterationBodyResult(
+                    DataStreamList.of(feedback), DataStreamList.of(feedback), terminationCriteria);
+        }
+    }
+
+    private static class MapWithState extends AbstractStreamOperator<Long>
+            implements OneInputStreamOperator<Long, Long> {
+        private ListState<Long> listState;
+        private ListState<Long> unionState;
+        private BroadcastState<Long, Long> broadcastState;
+
+        @Override
+        public void processElement(StreamRecord<Long> element) throws Exception {
+            long val = element.getValue();
+            listState.add(val);
+            unionState.add(val);
+            broadcastState.put(val, val);
+            output.collect(element);
+        }
+
+        @Override
+        public void initializeState(StateInitializationContext context) throws Exception {
+            super.initializeState(context);
+            listState =
+                    context.getOperatorStateStore()
+                            .getListState(new ListStateDescriptor<>("longState", Types.LONG));
+            unionState =
+                    context.getOperatorStateStore()
+                            .getUnionListState(new ListStateDescriptor<>("unionState", Types.LONG));
+            broadcastState =
+                    context.getOperatorStateStore()
+                            .getBroadcastState(
+                                    new MapStateDescriptor<>(
+                                            "broadcastState", Types.LONG, Types.LONG));
+        }
+    }
+
     private static class LongSink implements SinkFunction<Long> {
         private final SharedReference<BlockingQueue<Long>> collectedLong;