diff --git a/packages/client/hmi-client/src/components/model/petrinet/model-configurations/tera-stratified-matrix-modal.vue b/packages/client/hmi-client/src/components/model/petrinet/model-configurations/tera-stratified-matrix-modal.vue
index 80e447c16c..0ed1bfcba2 100644
--- a/packages/client/hmi-client/src/components/model/petrinet/model-configurations/tera-stratified-matrix-modal.vue
+++ b/packages/client/hmi-client/src/components/model/petrinet/model-configurations/tera-stratified-matrix-modal.vue
@@ -48,7 +48,6 @@ const props = defineProps<{
 	mmtParams: MiraTemplateParams;
 	id: string;
 	stratifiedMatrixType: StratifiedMatrix;
-	openValueConfig: boolean;
 }>();
 
 const emit = defineEmits(['close-modal', 'update-cell-value']);
diff --git a/packages/client/hmi-client/src/components/model/petrinet/model-diagrams/tera-model-diagram.vue b/packages/client/hmi-client/src/components/model/petrinet/model-diagrams/tera-model-diagram.vue
index 850bfe15ba..be4abcddd8 100644
--- a/packages/client/hmi-client/src/components/model/petrinet/model-diagrams/tera-model-diagram.vue
+++ b/packages/client/hmi-client/src/components/model/petrinet/model-diagrams/tera-model-diagram.vue
@@ -72,13 +72,12 @@
 			/>
 			<Teleport to="body">
 				<tera-stratified-matrix-modal
-					v-if="openValueConfig && modelConfiguration"
+					v-if="selectedTransitionId"
 					:id="selectedTransitionId"
 					:mmt="mmt"
 					:mmt-params="mmtParams"
 					:stratified-matrix-type="StratifiedMatrix.Rates"
-					:open-value-config="openValueConfig"
-					@close-modal="openValueConfig = false"
+					@close-modal="selectedTransitionId = ''"
 					@update-configuration="
 						(configToUpdate: ModelConfigurationLegacy) =>
 							emit('update-configuration', configToUpdate)
@@ -140,7 +139,6 @@ const isCollapsed = ref(true);
 const graphElement = ref<HTMLDivElement | null>(null);
 const graphLegendLabels = ref<string[]>([]);
 const graphLegendColors = ref<string[]>([]);
-const openValueConfig = ref(false);
 const selectedTransitionId = ref('');
 const modelType = computed(() => getModelType(props.model));
 const mmt = ref<MiraModel>(emptyMiraModel());
@@ -183,7 +181,6 @@ async function renderGraph() {
 		const { id, data } = selection.datum();
 		if (data.type === NodeType.Transition && data.isStratified) {
 			selectedTransitionId.value = id;
-			openValueConfig.value = true;
 		}
 	});
 
@@ -357,6 +354,10 @@ kbd {
 }
 
 .legend {
+	background-color: var(--surface-transparent);
+	backdrop-filter: blur(4px);
+	padding: var(--gap-xsmall) var(--gap-small);
+	border-radius: var(--border-radius);
 	position: absolute;
 	bottom: 0;
 	left: 0;
diff --git a/packages/client/hmi-client/src/model-representation/petrinet/nested-petrinet-renderer.ts b/packages/client/hmi-client/src/model-representation/petrinet/nested-petrinet-renderer.ts
index 84dad6b7fb..7610d07fa9 100644
--- a/packages/client/hmi-client/src/model-representation/petrinet/nested-petrinet-renderer.ts
+++ b/packages/client/hmi-client/src/model-representation/petrinet/nested-petrinet-renderer.ts
@@ -74,12 +74,10 @@ export class NestedPetrinetRenderer extends PetrinetRenderer {
 			d.matrixRows = matrixRowLen;
 			d.matrixCols = matrixColLen;
 
-			// FIXME: Consider rendering 1x1 matrices as a regular transition instead
-			d.data.isStratified = true;
+			if (matrixRowLen > 1 || matrixColLen > 1) d.data.isStratified = true;
 
 			// Initialize aspectRatio to 1 in case the matrix is square or empty
 			d.aspectRatio = 1;
-
 			// Check and set the aspect ratio based on the dimensions of the matrix
 			if (matrixRowLen > matrixColLen) {
 				d.aspectRatio = matrixColLen / matrixRowLen;
@@ -99,10 +97,77 @@ export class NestedPetrinetRenderer extends PetrinetRenderer {
 		});
 
 		const species = selection.filter((d) => d.data.type === NodeType.State);
-		const transitions = selection.filter((d) => d.data.type === NodeType.Transition);
+		const stratifiedTransitions = selection.filter(
+			(d) => d.data.type === NodeType.Transition && d.data.isStratified === true
+		);
+		const transitions = selection.filter(
+			(d) => d.data.type === NodeType.Transition && !d.data.isStratified
+		);
+
+		// species
+		species
+			.append('circle')
+			.classed('shape selectableNode', true)
+			.attr('r', (d) => 0.55 * d.width) // FIXME: need to adjust edge from sqaure mapping to circle
+			.attr('fill', (d) =>
+				d.data.strataType ? getNodeTypeColor(d.data.strataType) : getNestedTypeColor('base')
+			)
+			.attr('stroke', 'var(--petri-nodeBorder)')
+			.attr('stroke-width', 1)
+			.style('cursor', 'pointer');
 
 		// transitions
 		transitions
+			.append('rect')
+			.classed('shape selectableNode', true)
+			.attr('width', (d) => d.width)
+			.attr('height', (d) => d.height)
+			.attr('y', (d) => -d.height * 0.5)
+			.attr('x', (d) => -d.width * 0.5)
+			.attr('rx', '6')
+			.attr('ry', '6')
+			.style('fill', (d) =>
+				d.data.strataType ? getNodeTypeColor(d.data.strataType) : 'var(--petri-nodeFill'
+			)
+			.style('cursor', 'pointer')
+			.attr('stroke', 'var(--petri-nodeBorder)')
+			.attr('stroke-width', 1);
+
+		// transitions label text
+		transitions
+			.append('text')
+			.attr('y', (d) => setFontSize(d.id) / 4)
+			.style('text-anchor', 'middle')
+			.classed('latex-font', true)
+			.style('font-style', 'italic')
+			.style('font-size', (d) => setFontSize(d.id))
+			.style('stroke', '#FFF')
+			.style('paint-order', 'stroke')
+			.style('fill', 'var(--text-color-primary')
+			.style('pointer-events', 'none')
+			.html((d) => d.id);
+
+		// transitions expression text
+		transitions
+			.append('text')
+			.attr('y', (d) => -d.height / 2 - 8)
+			.classed('latex-fontt', true)
+			.style('font-style', 'italic')
+			.style('font-size', FONT_SIZE_SMALL)
+			.style('text-anchor', 'middle')
+			.style('paint-order', 'stroke')
+			.style('stroke', '#FFF')
+			.style('stroke-width', '3px')
+			.style('stroke-linecap', 'butt')
+			.style('fill', 'var(--text-color-primary')
+			.style('pointer-events', 'none')
+			.html((d) => {
+				if (d.data.expression) return d.data.expression;
+				return '';
+			});
+
+		// stratified transitions
+		stratifiedTransitions
 			.append('rect')
 			.classed('shape selectableNode', true)
 			.attr('width', (d) => ((d.aspectRatio ?? 1) >= 1 ? d.width : d.width))
@@ -116,18 +181,6 @@ export class NestedPetrinetRenderer extends PetrinetRenderer {
 			.attr('stroke', 'var(--petri-nodeBorder)')
 			.attr('stroke-width', 1);
 
-		// species
-		species
-			.append('circle')
-			.classed('shape selectableNode', true)
-			.attr('r', (d) => 0.55 * d.width) // FIXME: need to adjust edge from sqaure mapping to circle
-			.attr('fill', (d) =>
-				d.data.strataType ? getNodeTypeColor(d.data.strataType) : getNestedTypeColor('base')
-			)
-			.attr('stroke', 'var(--petri-nodeBorder)')
-			.attr('stroke-width', 1)
-			.style('cursor', 'pointer');
-
 		const renderNestedNodes = (
 			node: { [baseNodeId: string]: any },
 			parentRadius: number,
@@ -176,7 +229,7 @@ export class NestedPetrinetRenderer extends PetrinetRenderer {
 			renderNestedNodes(nestedMap, parentRadius, 0, 0, g, idx, 1);
 		});
 
-		transitions.each((d, idx, g) => {
+		stratifiedTransitions.each((d, idx, g) => {
 			const transitionMatrix = this.transitionMatrices?.[d.id] ?? [];
 
 			const matrixRowLen = transitionMatrix.length;
@@ -201,63 +254,10 @@ export class NestedPetrinetRenderer extends PetrinetRenderer {
 							.attr('stroke', '#ffffff')
 							.attr('stroke-width', 1);
 					}
-					// Draw label for number of columns
-					// transitionNode
-					// 	.append('text')
-					// 	.attr('x', 0)
-					// 	.attr('y', -d.height * 0.6)
-					// 	.attr('text-anchor', 'middle') // This will center-align the text horizontally
-					// 	.text(matrixColLen)
-					// 	.style('fill', '#cccccc')
-					// 	.style('font-size', '7px');
-
-					// Draw label for number of rows
-					// transitionNode
-					// 	.append('text')
-					// 	.attr('x', (-d.width * d.aspectRatio!) / 2 - 8)
-					// 	.attr('y', (-d.height * d.aspectRatio!) / 2 + 12)
-					// 	.attr('text-anchor', 'right') // This will center-align the text horizontally
-					// 	.text(matrixRowLen)
-					// 	.style('fill', '#cccccc')
-					// 	.style('font-size', '7px');
 				});
 			});
 		});
 
-		/* Don't show transition labels because we're showing matrices here */
-		// transitions label text
-		// transitions
-		// 	.append('text')
-		// 	.attr('y', () => 5)
-		// 	.style('text-anchor', 'middle')
-		// 	.style('paint-order', 'stroke')
-		// 	.style('fill', 'var(--text-color-primary')
-		// 	.style('pointer-events', 'none')
-		// 	.html((d) => d.id);
-
-		// transitions expression text
-		transitions
-			.append('text')
-			.attr('y', (d) => -d.height / 2 - 8)
-			.classed('latex-font', true)
-			.style('font-style', 'italic')
-			.style('font-size', FONT_SIZE_SMALL)
-			.style('text-anchor', 'middle')
-			.style('paint-order', 'stroke')
-			.style('stroke', '#FFF')
-			.style('stroke-width', '3px')
-			.style('stroke-linecap', 'butt')
-			.style('fill', 'var(--text-color-primary')
-			.style('pointer-events', 'none')
-			.html((d) => {
-				if (!this.graph.amr) return '';
-				const rate = this.graph.amr.semantics.ode?.rates?.find((r) => r.target === d.id);
-				if (rate) {
-					return rate.expression;
-				}
-				return '';
-			});
-
 		// species text
 		species
 			.append('text')